Commit b8435e63 authored by evshiron's avatar evshiron
Browse files

add --cors-allow-origins cmd opt

parent 89722fb5
Loading
Loading
Loading
Loading
+4 −3
Original line number Diff line number Diff line
@@ -86,6 +86,7 @@ parser.add_argument("--nowebui", action='store_true', help="use api=True to laun
parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI")
parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None)
parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False)
parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origins", default=None)

cmd_opts = parser.parse_args()
restricted_opts = {
+9 −0
Original line number Diff line number Diff line
@@ -5,6 +5,7 @@ import importlib
import signal
import threading
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware

from modules.paths import script_path
@@ -93,6 +94,11 @@ def initialize():
    signal.signal(signal.SIGINT, sigint_handler)


def setup_cors(app):
    if cmd_opts.cors_allow_origins:
        app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_allow_origins.split(','), allow_methods=['*'])


def create_api(app):
    from modules.api.api import Api
    api = Api(app, queue_lock)
@@ -114,6 +120,7 @@ def api_only():
    initialize()

    app = FastAPI()
    setup_cors(app)
    app.add_middleware(GZipMiddleware, minimum_size=1000)
    api = create_api(app)

@@ -147,6 +154,8 @@ def webui():
        # runnnig its code. We disable this here. Suggested by RyotaK.
        app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']

        setup_cors(app)

        app.add_middleware(GZipMiddleware, minimum_size=1000)

        if launch_api: