Update main.py

main.py

@@ -50,15 +50,37 @@ def install_dependents(skip: bool = False):
         return
 
     torch_index_url = os.environ.get("TORCH_INDEX_URL", "https://download.pytorch.org/whl/cu121")
+    logger.std_info(f"[Fooocus-API] Using torch index URL: {torch_index_url}")
+
+    # Check if you need pip install
     if not requirements_check():
+        logger.std_info("[Fooocus-API] Installing requirements.txt...")
         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")
+        logger.std_info("[Fooocus-API] Installing PyTorch with CUDA support...")
+        run_pip(
+            f"install torch==2.1.0 torchvision==0.16.0 --extra-index-url {torch_index_url}",
+            desc="torch",
+        )
+    else:
+        logger.std_info("[Fooocus-API] PyTorch with CUDA already installed")
 
 def preload_pipeline():
-    """Preload pipeline"""
+    """Preload pipeline with detailed error handling"""
     logger.std_info("[Fooocus-API] Preloading pipeline ...")
-    import modules.default_pipeline as _
+    try:
+        import torch
+        logger.std_info(f"[Fooocus-API] PyTorch version: {torch.__version__}, CUDA available: {torch.cuda.is_available()}")
+        if torch.cuda.is_available():
+            logger.std_info(f"[Fooocus-API] CUDA device: {torch.cuda.current_device()}, {torch.cuda.get_device_name(0)}")
+        
+        import modules.default_pipeline as pipeline
+        logger.std_info("[Fooocus-API] Pipeline module imported successfully")
+        # Add more granular steps here if needed to isolate crash
+    except Exception as e:
+        logger.std_error(f"[Fooocus-API] Pipeline preload failed: {str(e)}")
+        raise
 
 def prepare_environments(args) -> bool:
     """
@@ -69,11 +91,21 @@ def prepare_environments(args) -> bool:
     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"  # For base_url display, use localhost
+            host = "127.0.0.1"  # For base_url display
         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"))
+    logger.std_info(f"[Fooocus-API] Origin preset folder: {origin_preset_folder}")
+    logger.std_info(f"[Fooocus-API] Local preset folder: {preset_folder}")
+    # 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
@@ -90,8 +122,11 @@ def prepare_environments(args) -> bool:
     default.available_aspect_ratios = [default.get_aspect_ratio_value(a) for a in config.available_aspect_ratios]
 
     if not args.disable_preset_download:
+        logger.std_info("[Fooocus-API] Downloading models...")
         download_models()
+        logger.std_info("[Fooocus-API] Model download completed")
 
+    # Init task queue
     from fooocusapi import worker
     from fooocusapi.task_queue import TaskQueue
 
@@ -110,8 +145,13 @@ def prepare_environments(args) -> bool:
     return True
 
 def pre_setup():
-    """Pre setup, for replicate"""
+    """
+    Pre setup, for replicate or Hugging Face Spaces
+    """
     class Args(object):
+        """
+        Arguments object
+        """
         host = "127.0.0.1"
         port = 7860
         base_url = None
@@ -129,7 +169,7 @@ def pre_setup():
         gpu_device_id = None
         apikey = None
 
-    print("[Pre Setup] Prepare environments")
+    logger.std_info("[Pre Setup] Preparing environments")
     arguments = Args()
     sys.argv = [sys.argv[0]]
     sys.argv.append("--disable-image-log")
@@ -141,7 +181,7 @@ def pre_setup():
     task_thread = Thread(target=task_schedule_loop, daemon=True)
     task_thread.start()
 
-    print("[Pre Setup] Finished")
+    logger.std_info("[Pre Setup] Finished")
 
 if __name__ == "__main__":
     logger.std_info(f"[Fooocus API] Python {sys.version}")
@@ -161,7 +201,7 @@ if __name__ == "__main__":
     if prepare_environments(args):
         sys.argv = [sys.argv[0]]
 
-        # Load pipeline in new thread
+        # Load pipeline in new thread with error handling
         preload_pipeline_thread = Thread(target=preload_pipeline, daemon=True)
         preload_pipeline_thread.start()
 
@@ -170,26 +210,11 @@ if __name__ == "__main__":
         task_schedule_thread = Thread(target=task_schedule_loop, daemon=True)
         task_schedule_thread.start()
 
-        # Start API server with CORS
+        # Start API server using original Fooocus API call
         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}
-
-        # Mount the existing app routes
-        start_app(args, app=app)
+        try:
+            logger.std_info("[Fooocus-API] Starting API server...")
+            start_app(args)
+        except Exception as e:
+            logger.std_error(f"[Fooocus-API] Failed to start API server: {str(e)}")
+            raise