Update app.py

app.py

@@ -48,72 +48,109 @@ def add_extra_model_paths() -> None:
     """Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path."""
     try:
         from main import load_extra_path_config
+        extra_model_paths = find_path("extra_model_paths.yaml")
+        if extra_model_paths is not None:
+            load_extra_path_config(extra_model_paths)
+        else:
+            print("Could not find the extra_model_paths config file.")
     except ImportError:
-        print("Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead.")
-        from utils.extra_config import load_extra_path_config
-
-    extra_model_paths = find_path("extra_model_paths.yaml")
-    if extra_model_paths is not None:
-        load_extra_path_config(extra_model_paths)
-    else:
-        print("Could not find the extra_model_paths config file.")
+        try:
+            from utils.extra_config import load_extra_path_config
+            extra_model_paths = find_path("extra_model_paths.yaml")
+            if extra_model_paths is not None:
+                load_extra_path_config(extra_model_paths)
+            else:
+                print("Could not find the extra_model_paths config file.")
+        except ImportError:
+            print("Could not import extra config. Continuing without extra model paths.")
 
 add_comfyui_directory_to_sys_path()
-add_extra_model_paths()
+try:
+    add_extra_model_paths()
+except Exception as e:
+    print(f"Warning: Could not load extra model paths: {e}")
 
 def import_custom_nodes() -> None:
     """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS"""
-    import asyncio
-    import execution
-    from nodes import init_extra_nodes
-    import server
-
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
-    server_instance = server.PromptServer(loop)
-    execution.PromptQueue(server_instance)
-    init_extra_nodes()
+    try:
+        import asyncio
+        import execution
+        from nodes import init_extra_nodes
+        import server
+
+        # Check if we're already in an event loop
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # We're in an existing loop, use it
+                pass
+            else:
+                # Loop exists but not running, set a new one
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+        except RuntimeError:
+            # No loop exists, create one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        server_instance = server.PromptServer(loop)
+        execution.PromptQueue(server_instance)
+        init_extra_nodes()
+    except Exception as e:
+        print(f"Warning: Could not initialize custom nodes: {e}")
+        print("Continuing with basic ComfyUI nodes only...")
 
 from nodes import NODE_CLASS_MAPPINGS
 
 # Pre-load models outside the decorated function for ZeroGPU efficiency
-import_custom_nodes()
-
-# Initialize model loaders
-dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-dualcliploader_54 = dualcliploader.load_clip(
-    clip_name1="clip_l.safetensors",
-    clip_name2="t5xxl_fp16.safetensors",
-    type="flux",
-    device="default",
-)
-
-upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
-upscalemodelloader_44 = upscalemodelloader.load_model(model_name="4x-UltraSharp.pth")
-
-vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-vaeloader_55 = vaeloader.load_vae(vae_name="ae.safetensors")
-
-unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-unetloader_58 = unetloader.load_unet(
-    unet_name="flux1-dev.safetensors", weight_dtype="default"
-)
-
-downloadandloadflorence2model = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
-downloadandloadflorence2model_52 = downloadandloadflorence2model.loadmodel(
-    model="microsoft/Florence-2-large", precision="fp16", attention="sdpa"
-)
-
-# Pre-load models to GPU for efficiency
-from comfy import model_management
-model_loaders = [dualcliploader_54, vaeloader_55, unetloader_58, downloadandloadflorence2model_52]
-valid_models = [
-    getattr(loader[0], 'patcher', loader[0])
-    for loader in model_loaders
-    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
-]
-model_management.load_models_gpu(valid_models)
+try:
+    import_custom_nodes()
+
+    # Initialize model loaders
+    dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+    dualcliploader_54 = dualcliploader.load_clip(
+        clip_name1="clip_l.safetensors",
+        clip_name2="t5xxl_fp16.safetensors",
+        type="flux",
+        device="default",
+    )
+
+    upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
+    upscalemodelloader_44 = upscalemodelloader.load_model(model_name="4x-UltraSharp.pth")
+
+    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+    vaeloader_55 = vaeloader.load_vae(vae_name="ae.safetensors")
+
+    unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+    unetloader_58 = unetloader.load_unet(
+        unet_name="flux1-dev.safetensors", weight_dtype="default"
+    )
+
+    downloadandloadflorence2model = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
+    downloadandloadflorence2model_52 = downloadandloadflorence2model.loadmodel(
+        model="microsoft/Florence-2-large", precision="fp16", attention="sdpa"
+    )
+
+    # Pre-load models to GPU for efficiency
+    try:
+        from comfy import model_management
+        model_loaders = [dualcliploader_54, vaeloader_55, unetloader_58, downloadandloadflorence2model_52]
+        valid_models = [
+            getattr(loader[0], 'patcher', loader[0])
+            for loader in model_loaders
+            if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
+        ]
+        model_management.load_models_gpu(valid_models)
+        print("Models successfully pre-loaded to GPU")
+    except Exception as e:
+        print(f"Warning: Could not pre-load models to GPU: {e}")
+
+    print("ComfyUI setup completed successfully!")
+    
+except Exception as e:
+    print(f"Error during ComfyUI setup: {e}")
+    print("Please check that all required custom nodes are installed.")
+    raise
 
 @spaces.GPU(duration=120)  # Adjust duration based on your workflow speed
 def enhance_image(image_input, upscale_factor, steps, cfg_scale, denoise_strength, guidance_scale):