Update app.py

app.py

@@ -32,48 +32,56 @@ ip_adapter_path = "./"
 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 or download them with retry logic
-kolors_weights = model_path + "diffusion_pytorch_model.safetensors"
+# Check if base model weights exist or download them
+kolors_weights = model_path + "diffusers_weights.safetensors"
 if not os.path.exists(kolors_weights):
-    kolors_weights_unet = "./unet/diffusion_pytorch_model.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}")
-                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
-if not os.path.exists(ip_adapter_path + "ip-adapter.bin"):
+    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
+            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)
+                    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.")
+        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}")
 
 # Initialize model with empty weights
 with init_empty_weights():
     pipe = StableDiffusionXLPipeline.from_pretrained(
-        "Kwai-Kolors/Kolors-diffusers",
+        "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("h94/IP-Adapter-FaceID-Plus-SDXL", subfolder=None, weight_name="ip-adapter.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)
@@ -101,7 +109,7 @@ interface = gr.Interface(
     fn=generate_image,
     inputs=[gr.Image(type="pil", label="Upload Reference Image"), gr.Textbox(label="Enter Prompt", placeholder="e.g., A photorealistic astronaut in space")],
     outputs=[gr.Textbox(label="Status"), gr.Image(label="Generated Image")],
-    title="Face Reference Image Generator (Kolors with IP-Adapter)",
+    title="Face Reference Image Generator (Kolors-IP-Adapter-FaceID-Plus)",
     description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
 )