Update app.py

app.py

@@ -8,6 +8,7 @@ from huggingface_hub import hf_hub_download
 import os
 import logging
 from safetensors.torch import load_file
+import time
 
 # Set up detailed logging
 logging.basicConfig(level=logging.INFO)
@@ -34,24 +35,29 @@ except Exception as e:
     logger.error(f"Failed to load InsightFace model: {e}")
     raise
 
-# Download function with explicit path return
-def download_file(repo_id, filename, local_dir):
+# Download function with retry logic
+def download_file(repo_id, filename, local_dir, max_retries=3):
     file_path = os.path.join(local_dir, filename)
     if not os.path.exists(file_path):
-        logger.info(f"Downloading {filename} from {repo_id} to {local_dir}...")
-        try:
-            downloaded_path = hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                local_dir=local_dir,
-                cache_dir=cache_dir,
-                local_files_only=False
-            )
-            logger.info(f"Downloaded to {downloaded_path}")
-            return downloaded_path
-        except Exception as e:
-            logger.error(f"Download failed: {e}")
-            raise
+        for attempt in range(max_retries):
+            logger.info(f"Attempt {attempt + 1}/{max_retries}: Downloading {filename} from {repo_id} to {local_dir}...")
+            try:
+                downloaded_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename=filename,
+                    local_dir=local_dir,
+                    cache_dir=cache_dir,
+                    local_files_only=False
+                )
+                logger.info(f"Downloaded to {downloaded_path}")
+                return downloaded_path
+            except Exception as e:
+                logger.error(f"Download attempt {attempt + 1} failed: {e}")
+                if attempt < max_retries - 1:
+                    logger.info("Retrying in 5 seconds...")
+                    time.sleep(5)
+                else:
+                    raise RuntimeError(f"Failed to download {filename} after {max_retries} attempts: {e}")
     else:
         logger.info(f"Using cached file at {file_path}")
         return file_path
@@ -62,7 +68,7 @@ ip_adapter_path = "./"
 os.makedirs(kolors_unet_path, exist_ok=True)
 os.makedirs(ip_adapter_path, exist_ok=True)
 
-# Download weights
+# Download weights with retries
 logger.info("Starting weights download...")
 kolors_weights = download_file(
     "Kwai-Kolors/Kolors",
@@ -75,28 +81,39 @@ ip_adapter_weights = download_file(
     ip_adapter_path
 )
 
-# Load the pipeline
+# Load the pipeline with verbose logging and retry logic
 logger.info("Loading Stable Diffusion XL base model...")
 try:
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=dtype,
-        safety_checker=None,
-        local_files_only=False,
-        cache_dir=cache_dir,
-        variant="fp16",
-        use_safetensors=True
-    )
-    logger.info("SDXL base model loaded successfully.")
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Attempt {attempt + 1}/{max_retries}: Loading SDXL model...")
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-xl-base-1.0",
+                torch_dtype=dtype,
+                safety_checker=None,
+                local_files_only=False,
+                cache_dir=cache_dir,
+                variant="fp16",
+                use_safetensors=True
+            )
+            logger.info("SDXL base model loaded successfully.")
+            break
+        except Exception as e:
+            logger.error(f"Load attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                logger.info("Retrying in 5 seconds...")
+                time.sleep(5)
+            else:
+                raise RuntimeError(f"Failed to load SDXL model after {max_retries} attempts: {e}")
 except Exception as e:
     logger.error(f"Failed to load SDXL base model: {e}")
     raise
 
-# Load Kolors unet weights with strict=False to ignore mismatches
+# Load Kolors unet weights with strict=False
 logger.info(f"Loading Kolors unet weights from {kolors_weights}...")
 try:
     state_dict = load_file(kolors_weights, device=device)
-    # Load with strict=False to ignore unexpected keys and size mismatches
     pipe.unet.load_state_dict(state_dict, strict=False)
     logger.info("Kolors unet weights loaded successfully (with ignored mismatches).")
 except Exception as e: