Update app.py

app.py

@@ -6,7 +6,7 @@ from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import os
-import urllib.request
+from huggingface_hub import hf_hub_download
 import time
 
 # Allow network access for runtime download
@@ -42,34 +42,51 @@ if not os.path.exists(kolors_weights):
             print("Preloading failed. Attempting runtime download with retry...")
             os.makedirs("./unet", exist_ok=True)
             max_retries = 3
-            correct_url = "https://huggingface.co/Kwai-Kolors/Kolors/raw/main/unet/diffusion_pytorch_model.fp16.safetensors"
             for attempt in range(max_retries):
                 try:
                     print(f"Download attempt {attempt + 1} of {max_retries}")
-                    urllib.request.urlretrieve(correct_url, kolors_weights_unet)
+                    hf_hub_download(
+                        repo_id="Kwai-Kolors/Kolors",
+                        filename="unet/diffusion_pytorch_model.fp16.safetensors",
+                        local_dir="./unet",
+                        local_files_only=False
+                    )
                     print("Kolors base weights downloaded to", kolors_weights_unet)
                     model_path = "./unet/"
                     kolors_weights = kolors_weights_unet
                     break
-                except urllib.error.HTTPError as e:
-                    print(f"Download attempt {attempt + 1} failed: HTTP Error {e.code} - {e.reason}")
-                    if attempt < max_retries - 1:
-                        time.sleep(5)
-                    else:
-                        raise FileNotFoundError(f"Failed to download Kolors base weights after {max_retries} attempts: HTTP Error {e.code} - {e.reason}. Verify the URL or contact support.")
                 except Exception as e:
                     print(f"Download attempt {attempt + 1} failed: {e}")
                     if attempt < max_retries - 1:
                         time.sleep(5)
                     else:
-                        raise FileNotFoundError(f"Failed to download Kolors base weights after {max_retries} attempts: {e}. Check network access or contact support.")
+                        raise FileNotFoundError(f"Failed to download Kolors base weights after {max_retries} attempts: {e}. Verify the repo or contact support.")
         else:
             model_path = "./unet/"
             kolors_weights = kolors_weights_unet
 
-# Check if IP-Adapter weights exist (preloaded)
-if not os.path.exists(ip_adapter_path + "ipa-faceid-plus.bin"):
-    raise FileNotFoundError(f"IP-Adapter weights not found at {ip_adapter_path}")
+# Check if IP-Adapter weights exist or download them
+ip_adapter_weights = ip_adapter_path + "ipa-faceid-plus.bin"
+if not os.path.exists(ip_adapter_weights):
+    print("IP-Adapter preloading failed. Attempting runtime download with retry...")
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            print(f"IP-Adapter download attempt {attempt + 1} of {max_retries}")
+            hf_hub_download(
+                repo_id="Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
+                filename="ipa-faceid-plus.bin",
+                local_dir="./",
+                local_files_only=False
+            )
+            print("IP-Adapter weights downloaded to", ip_adapter_weights)
+            break
+        except Exception as e:
+            print(f"IP-Adapter download attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                time.sleep(5)
+            else:
+                raise FileNotFoundError(f"Failed to download IP-Adapter weights after {max_retries} attempts: {e}. Verify the repo or contact support.")
 
 # Initialize model with empty weights
 with init_empty_weights():