Update app.py

app.py

@@ -5,6 +5,24 @@ from diffusers import DiffusionPipeline
 import torch
 import time
 import psutil
+from huggingface_hub import snapshot_download # New import
+
+# ... (Other imports and variables remain the same)
+
+def load_pipeline(model_id):
+    if model_id in PIPELINES:
+        return PIPELINES[model_id]
+    else:
+        # Download model using snapshot_download to handle LFS files
+        model_path = snapshot_download(repo_id=model_id, local_dir="./models")  # Download and cache models
+        pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32) 
+        pipe.to(DEVICE)
+        PIPELINES[model_id] = pipe
+        return pipe
+
+
+# ... (Rest of your code remains the same)
+
 
 # Get the number of physical CPU cores (excluding hyperthreads)
 NUM_CPU_CORES = psutil.cpu_count(logical=True)