updating app.py

app.py

@@ -6,6 +6,7 @@ import torch
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import spaces
+from comfy import model_management
 
 hf_hub_download(repo_id="Comfy-Org/stable-diffusion-v1-5-archive", filename="v1-5-pruned-emaonly-fp16.safetensors", local_dir="models/checkpoints")
 
@@ -119,11 +120,30 @@ def import_custom_nodes() -> None:
 
 from nodes import NODE_CLASS_MAPPINGS
 
+checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
+checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
+    ckpt_name="v1-5-pruned.safetensors"
+)
+
+
+#Add all the models that load a safetensors file
+model_loaders = [checkpointloadersimple_4]
+
+# Check which models are valid and how to best load them
+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)
+]
+
+#Finally loads the models
+model_management.load_models_gpu(valid_models)
+
 @spaces.GPU(duration=60) #modify the duration for the average it takes for your worflow to run, in seconds
 def generate_image(prompt):
     import_custom_nodes()
     with torch.inference_mode():
-        checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
+
         checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
             ckpt_name="v1-5-pruned.safetensors"
         )