import os
import gradio as gr
import torch
import numpy as np
from PIL import Image
import cv2

from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel
import insightface

# ========== CONFIG ==========
device = "cpu"
dtype = torch.float32

BASE_MODEL = "runwayml/stable-diffusion-v1-5"
CONTROLNET_MODEL = "lllyasviel/control_v11e_sd15_reference_only"

# ========== UTILS ==========
def pil_to_cv(img):
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def cv_to_pil(img):
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

def make_grid(images, rows, cols):
    w, h = images[0].size
    grid = Image.new("RGB", (cols*w, rows*h))
    for i, im in enumerate(images):
        grid.paste(im, box=((i % cols) * w, (i // cols) * h))
    return grid

# ========== LOAD MODELS ==========
print("🔄 Loading ControlNet...")
controlnet = ControlNetModel.from_pretrained(CONTROLNET_MODEL, torch_dtype=dtype)

print("🔄 Loading Stable Diffusion pipeline...")
pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnet,
    torch_dtype=dtype,
    safety_checker=None
).to(device)

print("🔄 Loading InsightFace (antelopev2)...")
ins_model = insightface.app.FaceAnalysis(name="antelopev2", providers=["CPUExecutionProvider"])
ins_model.prepare(ctx_id=0, det_size=(640, 640))

# ========== MAIN FUNCTION ==========
def generate_grid(init_image, prompt, neg_prompt, steps, cfg, strength, seed, width, height):
    if init_image is None:
        return None, "❌ Upload foto dulu."

    # Deteksi wajah
    cv_img = pil_to_cv(init_image)
    faces = ins_model.get(cv_img)
    if not faces:
        return None, "❌ Wajah tidak terdeteksi."

    # Ambil wajah terbesar
    faces.sort(key=lambda f: (f.bbox[2]-f.bbox[0])*(f.bbox[3]-f.bbox[1]), reverse=True)
    x1, y1, x2, y2 = map(int, faces[0].bbox)
    face_crop = init_image.crop((x1, y1, x2, y2)).resize((width, height))

    # Generator seed
    g = torch.Generator(device=device)
    if seed >= 0:
        g.manual_seed(seed)

    images = []
    for i in range(9):
        out = pipe(
            prompt=prompt,
            negative_prompt=neg_prompt if neg_prompt.strip() else None,
            image=init_image.resize((width, height)),
            control_image=face_crop,
            num_inference_steps=steps,
            guidance_scale=cfg,
            strength=strength,
            generator=torch.Generator(device=device).manual_seed(
                seed if seed >= 0 else torch.seed()
            )
        ).images[0]
        images.append(out)

    return make_grid(images, 3, 3), f"✅ Face bbox: {x1},{y1},{x2},{y2}"

# ========== GRADIO UI ==========
with gr.Blocks() as demo:
    gr.Markdown("## 👤 CPU Face Consistency 3×3 Generator (ControlNet Reference-Only + InsightFace)")

    with gr.Row():
        with gr.Column():
            init_image = gr.Image(label="Upload Foto Wajah", type="pil")
            prompt = gr.Textbox(label="Prompt", value="studio portrait, soft lighting, ultra detailed skin")
            neg_prompt = gr.Textbox(label="Negative Prompt", value="blurry, deformed, bad anatomy")
            steps = gr.Slider(8, 25, 14, step=1, label="Steps")
            cfg = gr.Slider(3, 12, 7.0, step=0.5, label="CFG Scale")
            strength = gr.Slider(0.3, 0.8, 0.55, step=0.05, label="Denoising Strength")
            width = gr.Slider(320, 512, 384, step=32, label="Width")
            height = gr.Slider(320, 512, 384, step=32, label="Height")
            seed = gr.Number(value=-1, precision=0, label="Seed (-1 = acak)")
            run = gr.Button("Generate 3×3")
        with gr.Column():
            out_img = gr.Image(label="Output Grid")
            log = gr.Markdown()

    run.click(
        fn=generate_grid,
        inputs=[init_image, prompt, neg_prompt, steps, cfg, strength, seed, width, height],
        outputs=[out_img, log]
    )

if __name__ == "__main__":
    demo.launch()