app.py

import spaces
import gradio as gr
import torch
import numpy as np
import cv2
import hashlib
import os
from PIL import Image
from diffusers import FluxPipeline
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model

# --- GLOBAL CONFIG ---
MODEL_ID = "black-forest-labs/FLUX.1-dev"
HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize models as None for ZeroGPU lazy loading
face_app = None
swapper = None
pipe = None

def load_models_on_gpu():
    """Initializes models only when GPU is allocated."""
    global face_app, swapper, pipe
    
    if face_app is None:
        # Use CPU for detection to save GPU VRAM for the main Flux pipe
        face_app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
        face_app.prepare(ctx_id=0, det_size=(640, 640))

    if swapper is None:
        model_file = 'inswapper_128.onnx'
        if os.path.exists(model_file):
            swapper = get_model(model_file, providers=['CPUExecutionProvider'])

    if pipe is None:
        pipe = FluxPipeline.from_pretrained(
            MODEL_ID, 
            torch_dtype=torch.bfloat16,
            token=HF_TOKEN
        )
        pipe.enable_model_cpu_offload()

def upscale_image(image):
    img = np.array(image)
    w, h = image.size
    # Simple high-quality Lanczos upscale
    upscaled = cv2.resize(img, (w*2, h*2), interpolation=cv2.INTER_LANCZOS4)
    gaussian_blur = cv2.GaussianBlur(upscaled, (0, 0), 3)
    sharpened = cv2.addWeighted(upscaled, 1.5, gaussian_blur, -0.5, 0)
    return Image.fromarray(sharpened)

@spaces.GPU(duration=150)
def generate_vton_final(face_image, body_type, height_ft):
    if face_image is None: 
        return None, "Please upload a face image."
    
    load_models_on_gpu()

    # 1. Face Analysis
    img_np = np.array(face_image)
    cv_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    faces = face_app.get(cv_img)
    if not faces: 
        return None, "No face detected in the upload."
    
    source_face = faces[0]
    gender = "man" if source_face.gender == 1 else "woman"

    # 2. Body Type Mapping (Adding 'fat' via 'plus-size' keyword for better AI results)
    body_map = {
        "slim": "slender and lean",
        "muscular": "athletic and muscular",
        "average": "average",
        "plus-size": "heavyset, plus-size, large frame" 
    }
    refined_body = body_map.get(body_type, body_type)

    # 3. Prompt Construction
    profile_seed = int(hashlib.md5(f"{gender}-{body_type}-{height_ft}".encode()).hexdigest(), 16) % (10**9)
    generator = torch.Generator("cuda").manual_seed(profile_seed)

    prompt = (
        f"A full body 8k professional photo of a {gender}, {refined_body} build, {height_ft}ft tall. "
        f"Standing in a relaxed, natural pose, facing the camera. "
        f"Wearing stylish casual clothing, clean studio background, sharp focus, cinematic lighting."
    )

    # 4. Generation
    gen_img = pipe(
        prompt=prompt,
        height=1024, width=768,
        guidance_scale=3.5,
        num_inference_steps=28,
        generator=generator
    ).images[0]

    # 5. Face Swap
    if swapper:
        res_np = np.array(gen_img)
        res_cv = cv2.cvtColor(res_np, cv2.COLOR_RGB2BGR)
        target_faces = face_app.get(res_cv)
        if target_faces:
            # Sort by area to find the primary subject
            target_faces = sorted(target_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=True)
            res_cv = swapper.get(res_cv, target_faces[0], source_face, paste_back=True)
            gen_img = Image.fromarray(cv2.cvtColor(res_cv, cv2.COLOR_BGR2RGB))

    return upscale_image(gen_img), f"Success | Seed: {profile_seed}"

# --- GRADIO INTERFACE ---
with gr.Blocks() as demo:
    gr.Markdown("# 💎 AI Virtual Model Engine")
    gr.Markdown("Upload a portrait to generate a custom full-body model.")
    
    with gr.Row():
        with gr.Column():
            face_in = gr.Image(type="pil", label="Step 1: Upload Face")
            body_in = gr.Radio(
                ["slim", "muscular", "average", "plus-size"], 
                value="average", 
                label="Step 2: Body Build"
            )
            h_in = gr.Slider(4.5, 7.5, value=5.8, step=0.1, label="Step 3: Height (ft)")
            btn = gr.Button("Generate High-Res Model", variant="primary")
            
        with gr.Column():
            img_out = gr.Image(label="Final Result")
            status = gr.Textbox(label="Status")

    btn.click(generate_vton_final, [face_in, body_in, h_in], [img_out, status])

# Fixed Gradio 6.0 warning: theme moved to launch()
if __name__ == "__main__":
    demo.launch(ssr_mode=False)