Update app.py

app.py

@@ -6,9 +6,11 @@ from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import os
+import urllib.request
+import time
 
-# Force offline mode for Hugging Face Hub
-os.environ["HF_HUB_OFFLINE"] = "1"
+# Allow network access for runtime download
+os.environ["HF_HUB_OFFLINE"] = "0"
 
 # Set device to CPU
 device = "cpu"
@@ -22,24 +24,42 @@ try:
 except Exception as e:
     raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access for initial download.")
 
-# Define paths for preloaded weights
-model_path = "./"  # Start with root, adjust if needed
+# Define paths for preloaded or downloaded weights
+model_path = "./"  # Start with root
 ip_adapter_path = "./"
 
-# Debug: List files to confirm preloading
+# Debug: List files to confirm preloading or download
 print("Files in root directory:", os.listdir("."))
 print("Files in ./unet/ directory:", os.listdir("./unet") if os.path.exists("./unet") else "No ./unet/ directory")
 
-# Check if weights exist
+# Check if weights exist or download them with retry logic
 kolors_weights = model_path + "diffusion_pytorch_model.safetensors"
 if not os.path.exists(kolors_weights):
-    # Check if it's in ./unet/ as a fallback
     kolors_weights_unet = "./unet/diffusion_pytorch_model.safetensors"
-    if os.path.exists(kolors_weights_unet):
+    if not os.path.exists(kolors_weights_unet):
+        print("Preloading failed. Attempting runtime download with retry...")
+        os.makedirs("./unet", exist_ok=True)
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                print(f"Download attempt {attempt + 1} of {max_retries}")
+                urllib.request.urlretrieve(
+                    "https://huggingface.co/Kwai-Kolors/Kolors-diffusers/resolve/main/unet/diffusion_pytorch_model.safetensors",
+                    kolors_weights_unet
+                )
+                print("Kolors weights downloaded to", kolors_weights_unet)
+                model_path = "./unet/"
+                kolors_weights = kolors_weights_unet
+                break
+            except Exception as e:
+                print(f"Download attempt {attempt + 1} failed: {e}")
+                if attempt < max_retries - 1:
+                    time.sleep(5)  # Wait 5 seconds before retrying
+                else:
+                    raise FileNotFoundError(f"Failed to download Kolors weights after {max_retries} attempts: {e}. Check network access or contact support.")
+    else:
         model_path = "./unet/"
         kolors_weights = kolors_weights_unet
-    else:
-        raise FileNotFoundError(f"Kolors model weights not found at {kolors_weights} or {kolors_weights_unet}. Preloading failed. Please check .env configuration.")
 if not os.path.exists(ip_adapter_path + "ip-adapter.bin"):
     raise FileNotFoundError(f"IP-Adapter weights not found at {ip_adapter_path}")