Update app.py

app.py

@@ -25,7 +25,7 @@ except Exception as e:
     raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access for initial download.")
 
 # Define paths for preloaded or downloaded weights
-model_path = "./unet/"  # Updated to your unet directory
+model_path = "./"  # Start with root
 ip_adapter_path = "./"
 
 # Debug: List files to confirm preloading or download
@@ -33,39 +33,48 @@ 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 base model weights exist or download them
-kolors_weights = model_path + "diffusion_pytorch_model.fp16.safetensors"
+kolors_weights = model_path + "diffusers_weights.safetensors"
 if not os.path.exists(kolors_weights):
-    print("Kolors base weights not found. Attempting runtime download...")
-    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}")
-            hf_hub_download(
-                repo_id="AkashKumarave/my3",
-                filename="unet/diffusion_pytorch_model.fp16.safetensors",
-                local_dir="./unet",
-                local_files_only=False
-            )
-            print("Kolors base weights downloaded to", kolors_weights)
-            break
-        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}. Verify the repo or contact support.")
+    kolors_weights = model_path + "diffusion_pytorch_model.fp16.safetensors"
+    if not os.path.exists(kolors_weights):
+        kolors_weights_unet = "./unet/diffusion_pytorch_model.fp16.safetensors"
+        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}")
+                    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 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}. Verify the repo or contact support.")
+        else:
+            model_path = "./unet/"
+            kolors_weights = kolors_weights_unet
 
 # 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...")
+    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="AkashKumarave/my3",
+                repo_id="Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
                 filename="ipa-faceid-plus.bin",
                 local_dir="./",
                 local_files_only=False
@@ -82,14 +91,14 @@ if not os.path.exists(ip_adapter_weights):
 # Initialize model with empty weights
 with init_empty_weights():
     pipe = StableDiffusionXLPipeline.from_pretrained(
-        "AkashKumarave/my3",
+        "Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
         torch_dtype=dtype,
         safety_checker=None,
     )
 
 # Load and dispatch model with accelerate
 pipe = load_checkpoint_and_dispatch(pipe, model_path, device_map="cpu", offload_folder=None)
-pipe.load_ip_adapter("AkashKumarave/my3", subfolder=None, weight_name="ipa-faceid-plus.bin")
+pipe.load_ip_adapter("Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus", subfolder=None, weight_name="ipa-faceid-plus.bin")
 
 def generate_image(uploaded_image, prompt):
     img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
@@ -121,4 +130,4 @@ interface = gr.Interface(
     description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
 )
 
-interface.launch()
+interface.launch()