import cv2
import torch
import numpy as np
import gradio as gr
from diffusers import StableDiffusionXLPipeline
from insightface.app import FaceAnalysis
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
import os
import urllib.request

# Force offline mode for Hugging Face Hub (but allow InsightFace download)
os.environ["HF_HUB_OFFLINE"] = "1"

# Set device to CPU
device = "cpu"
dtype = torch.float32

# Set up InsightFace model directory
insightface_model_dir = "/home/user/.insightface/models/buffalo_l"
os.makedirs(insightface_model_dir, exist_ok=True)
buffalo_l_zip = "./buffalo_l.zip"

# Download buffalo_l.zip if not present
if not os.path.exists(buffalo_l_zip):
    try:
        print("Downloading buffalo_l.zip for InsightFace...")
        urllib.request.urlretrieve(
            "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip",
            buffalo_l_zip
        )
        print("Download completed.")
    except Exception as e:
        print(f"Failed to download buffalo_l.zip: {e}")
        raise RuntimeError("Cannot download buffalo_l.zip. Please ensure network access.")

# Extract buffalo_l.zip
if os.path.exists(buffalo_l_zip):
    import zipfile
    with zipfile.ZipFile(buffalo_l_zip, "r") as zip_ref:
        zip_ref.extractall(insightface_model_dir)
    print("Extracted buffalo_l.zip.")

# Load face encoder
face_app = FaceAnalysis(providers=["CPUExecutionProvider"], root="/home/user/.insightface/models")
face_app.prepare(ctx_id=0, det_size=(480, 480))

# Define paths for preloaded weights
model_path = "./unet/"  # Adjusted for preloaded Kolors weights
ip_adapter_path = "./"

# Check if Kolors weights exist, download if missing
kolors_weights = model_path + "diffusion_pytorch_model.safetensors"
if not os.path.exists(model_path):
    os.makedirs(model_path, exist_ok=True)
if not os.path.exists(kolors_weights):
    try:
        print("Downloading Kolors model weights...")
        urllib.request.urlretrieve(
            "https://huggingface.co/Kwai-Kolors/Kolors-diffusers/resolve/main/unet/diffusion_pytorch_model.safetensors",
            kolors_weights
        )
        print("Kolors weights downloaded.")
    except Exception as e:
        print(f"Failed to download Kolors weights: {e}")
        raise RuntimeError("Cannot download Kolors weights. Please ensure network access.")

# Check if IP-Adapter weights exist
if not os.path.exists(ip_adapter_path + "ip-adapter.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",
        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")

def generate_image(uploaded_image, prompt):
    img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
    faces = face_app.get(img)
    if not faces:
        return "No face detected!", None

    face_info = faces[-1]
    face_emb = face_info["embedding"]

    try:
        image = pipe(
            prompt=prompt,
            image_embeds=face_emb,
            num_inference_steps=20,
            guidance_scale=7.5,
            height=512,