Update main.py

main.py

@@ -40,7 +40,6 @@ default_command_live = True
 index_url = os.environ.get("INDEX_URL", "")
 re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
 
-
 def install_dependents(skip: bool = False):
     """
     Check and install dependencies
@@ -50,50 +49,31 @@ def install_dependents(skip: bool = False):
     if skip:
         return
 
-    torch_index_url = os.environ.get(
-        "TORCH_INDEX_URL", "https://download.pytorch.org/whl/cu121"
-    )
-
-    # Check if you need pip install
+    torch_index_url = os.environ.get("TORCH_INDEX_URL", "https://download.pytorch.org/whl/cu121")
     if not requirements_check():
         run_pip("install -r requirements.txt", "requirements")
-
     if not check_torch_cuda():
-        run_pip(
-            f"install torch==2.1.0 torchvision==0.16.0 --extra-index-url {torch_index_url}",
-            desc="torch",
-        )
-
+        run_pip(f"install torch==2.1.0 torchvision==0.16.0 --extra-index-url {torch_index_url}", desc="torch")
 
 def preload_pipeline():
     """Preload pipeline"""
     logger.std_info("[Fooocus-API] Preloading pipeline ...")
     import modules.default_pipeline as _
 
-
 def prepare_environments(args) -> bool:
     """
     Prepare environments
     Args:
         args: command line arguments
     """
-
     if args.base_url is None or len(args.base_url.strip()) == 0:
         host = args.host
         if host == "0.0.0.0":
-            host = "127.0.0.1"
+            host = "127.0.0.1"  # For base_url display, use localhost
         args.base_url = f"http://{host}:{args.port}"
 
     sys.argv = [sys.argv[0]]
 
-    # Define preset folder paths but avoid runtime file operations
-    origin_preset_folder = os.path.abspath(os.path.join(module_path, "presets"))
-    preset_folder = os.path.abspath(os.path.join(script_path, "presets"))
-    # Comment out file operations to avoid permission issues on Hugging Face Spaces
-    # if os.path.exists(preset_folder):
-    #     shutil.rmtree(preset_folder)
-    # shutil.copytree(origin_preset_folder, preset_folder)
-
     from modules import config
     from fooocusapi.configs import default
     from fooocusapi.utils.model_loader import download_models
@@ -106,17 +86,12 @@ def prepare_environments(args) -> bool:
     default.default_loras = config.default_loras
     default.default_cfg_scale = config.default_cfg_scale
     default.default_prompt_negative = config.default_prompt_negative
-    default.default_aspect_ratio = default.get_aspect_ratio_value(
-        config.default_aspect_ratio
-    )
-    default.available_aspect_ratios = [
-        default.get_aspect_ratio_value(a) for a in config.available_aspect_ratios
-    ]
+    default.default_aspect_ratio = default.get_aspect_ratio_value(config.default_aspect_ratio)
+    default.available_aspect_ratios = [default.get_aspect_ratio_value(a) for a in config.available_aspect_ratios]
 
     if not args.disable_preset_download:
         download_models()
 
-    # Init task queue
     from fooocusapi import worker
     from fooocusapi.task_queue import TaskQueue
 
@@ -130,18 +105,13 @@ def prepare_environments(args) -> bool:
     logger.std_info(f"[Fooocus-API] Task queue size: {args.queue_size}")
     logger.std_info(f"[Fooocus-API] Queue history size: {args.queue_history}")
     logger.std_info(f"[Fooocus-API] Webhook url: {args.webhook_url}")
+    logger.std_info(f"[Fooocus-API] Base URL: {args.base_url}")
 
     return True
 
-
 def pre_setup():
-    """
-    Pre setup, for replicate
-    """
+    """Pre setup, for replicate"""
     class Args(object):
-        """
-        Arguments object
-        """
         host = "127.0.0.1"
         port = 7860
         base_url = None
@@ -160,24 +130,19 @@ def pre_setup():
         apikey = None
 
     print("[Pre Setup] Prepare environments")
-
     arguments = Args()
     sys.argv = [sys.argv[0]]
     sys.argv.append("--disable-image-log")
 
     install_dependents(arguments.skip_pip)
-
     prepare_environments(arguments)
 
-    # Start task schedule thread
     from fooocusapi.worker import task_schedule_loop
-
     task_thread = Thread(target=task_schedule_loop, daemon=True)
     task_thread.start()
 
     print("[Pre Setup] Finished")
 
-
 if __name__ == "__main__":
     logger.std_info(f"[Fooocus API] Python {sys.version}")
     logger.std_info(f"[Fooocus API] Fooocus API version: {version}")
@@ -186,6 +151,7 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser()
     add_base_args(parser, True)
+    parser.set_defaults(host="0.0.0.0")  # Default to 0.0.0.0 for broader access
 
     args, _ = parser.parse_known_args()
     install_dependents(skip=args.skip_pip)
@@ -201,11 +167,29 @@ if __name__ == "__main__":
 
         # Start task schedule thread
         from fooocusapi.worker import task_schedule_loop
-
         task_schedule_thread = Thread(target=task_schedule_loop, daemon=True)
         task_schedule_thread.start()
 
-        # Start api server
+        # Start API server with CORS
         from fooocusapi.api import start_app
+        from fastapi import FastAPI
+        from fastapi.middleware.cors import CORSMiddleware
+
+        app = FastAPI()
+
+        # Add CORS middleware
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],  # Allow all origins (adjust for production)
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+        # Add config endpoint to expose base URL
+        @app.get("/config")
+        async def get_config():
+            return {"base_url": args.base_url}
 
-        start_app(args)
+        # Mount the existing app routes
+        start_app(args, app=app)