Update app.py

app.py

@@ -2,15 +2,15 @@ import cv2
 import torch
 import numpy as np
 import gradio as gr
-from diffusers import StableDiffusionXLPipeline, ControlNetModel
+from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import os
 
-# Force offline mode to avoid runtime Hub connections
+# Force offline mode
 os.environ["HF_HUB_OFFLINE"] = "1"
 
-# Set device to CPU (free tier has no GPU)
+# Set device to CPU
 device = "cpu"
 dtype = torch.float32
 
@@ -19,29 +19,26 @@ face_app = FaceAnalysis(providers=["CPUExecutionProvider"])
 face_app.prepare(ctx_id=0, det_size=(480, 480))
 
 # Define paths for preloaded weights
-controlnet_path = "./ControlNetModel"
-face_adapter_path = "./ip-adapter.bin"
+model_path = "./"  # Kolors base model
+ip_adapter_path = "./ip-adapter-faceid-plus-sdxl"
 
 # Check if files exist
-if not os.path.exists(controlnet_path) or not os.path.exists(os.path.join(controlnet_path, "config.json")):
-    raise FileNotFoundError(f"ControlNetModel directory or config.json not found at {controlnet_path}")
-if not os.path.exists(face_adapter_path):
-    raise FileNotFoundError(f"ip-adapter.bin not found at {face_adapter_path}")
+if not os.path.exists(model_path + "diffusion_pytorch_model.safetensors"):
+    raise FileNotFoundError(f"Kolors model weights not found at {model_path}")
+if not os.path.exists(ip_adapter_path) or not os.path.exists(os.path.join(ip_adapter_path, "config.json")):
+    raise FileNotFoundError(f"IP-Adapter weights or config.json not found at {ip_adapter_path}")
 
-# Initialize models with empty weights
+# Initialize model with empty weights
 with init_empty_weights():
-    controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=dtype)
     pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        controlnet=controlnet,
+        "Kwai-Kolors/Kolors-diffusers",
         torch_dtype=dtype,
         safety_checker=None,
     )
 
-# Load and dispatch models with accelerate
-controlnet = load_checkpoint_and_dispatch(controlnet, controlnet_path, device_map="cpu", offload_folder=None)
-pipe = load_checkpoint_and_dispatch(pipe, "./", device_map="cpu", offload_folder=None)
-pipe.load_ip_adapter_instantid(face_adapter_path)
+# 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")
 
 def generate_image(uploaded_image, prompt):
     # Convert Gradio image to OpenCV format
@@ -61,7 +58,6 @@ def generate_image(uploaded_image, prompt):
             guidance_scale=7.5,
             height=512,
             width=512,
-            controlnet_conditioning_scale=1.0,
         ).images[0]
         return "Image generated successfully!", image
     except Exception as e:
@@ -78,7 +74,7 @@ interface = gr.Interface(
         gr.Textbox(label="Status"),
         gr.Image(label="Generated Image")
     ],
-    title="Face Reference Image Generator",
+    title="Face Reference Image Generator (Kolors with IP-Adapter)",
     description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
 )