Update app.py

app.py

@@ -7,6 +7,29 @@ import model_loader
 import pipeline
 from transformers import CLIPTokenizer
 from config import Config, DeviceConfig
+import os
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+
+# Create data directory if it doesn't exist
+data_dir = Path("data")
+data_dir.mkdir(exist_ok=True)
+
+# Model configuration
+MODEL_REPO = "stabilityai/stable-diffusion-v1-5"
+MODEL_FILENAME = "v1-5-pruned-emaonly.ckpt"
+model_file = data_dir / MODEL_FILENAME
+
+# Download model if it doesn't exist
+if not model_file.exists():
+    print(f"Downloading model from {MODEL_REPO}...")
+    model_file = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILENAME,
+        local_dir=data_dir,
+        local_dir_use_symlinks=False
+    )
+    print("Model downloaded successfully!")
 
 # Device configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -19,8 +42,7 @@ config = Config(
 )
 
 # Load models with SE blocks enabled
-model_file = "data/v1-5-pruned-emaonly.ckpt"
-config.models = model_loader.load_models(model_file, device, use_se=True)
+config.models = model_loader.load_models(str(model_file), device, use_se=True)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024