app.py

import spaces
import gradio as gr
import torch
import cv2
import numpy as np
import os
from PIL import Image, ImageDraw, ImageFilter
from diffusers import StableDiffusionXLInpaintPipeline
from insightface.app import FaceAnalysis
import insightface
from huggingface_hub import hf_hub_download, snapshot_download

# ==============================================================================
# 1. SETUP DELL'AMBIENTE E DOWNLOAD MODELLI
# ==============================================================================

# Download sicuro e robusto di antelopev2 da Hugging Face per evitare errori di unzip
print("Download modelli AntelopeV2 da Hugging Face...")
try:
    snapshot_download("DIAMONIK7777/antelopev2", local_dir="./models/antelopev2")
except Exception as e:
    print(f"Errore download AntelopeV2: {e}")

if not os.path.exists("./models/inswapper_128.onnx"):
    print("Download Inswapper...")
    hf_hub_download(repo_id="ezioruan/inswapper_128.onnx", filename="inswapper_128.onnx", local_dir="./models")

# Inizializza l'analizzatore di volti (InsightFace)
print("Inizializzazione FaceAnalysis...")
face_app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
face_app.prepare(ctx_id=0, det_size=(640, 640))

# Inizializza Inswapper (per fissare l'identità alla fine)
print("Inizializzazione Inswapper...")
swapper = insightface.model_zoo.get_model('./models/inswapper_128.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# Caricamento Pipeline Inpainting SDXL
print("Caricamento pipeline SDXL Inpaint...")
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else "float32"

pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
    torch_dtype=dtype,
    variant="fp16" if torch.cuda.is_available() else None,
    use_safetensors=True
)

# Aggiunta IP-Adapter per l'analisi semantica (Capelli, Occhiali, Barba)
print("Caricamento IP-Adapter per SDXL...")
pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
pipe.set_ip_adapter_scale(0.7) # Bilanciamento tra fedeltà al selfie e stile del prompt

pipe.watermark = None 

# ==============================================================================
# 2. CONFIGURAZIONE STILI E TEMPLATE
# ==============================================================================
STYLES = {
    "Stile 1: Bit-Builder (8-Bit/Pixel Art)": {
        "template": "esempi/8bit.png",
        "repo_id": "nerijs/pixel-art-xl",
        "weight_name": "pixel-art-xl.safetensors",
        "prompt": "pixelart style, 8-bit game style face and hair, perfect pixel art shading, 8-bit retro aesthetic, matching lighting",
        "negative_prompt": "ugly, deformed, lowres, realistic, photographic, 3d render, photo, photorealistic"
    },
    "Stile 2: Il Costruttore Rinascimentale": {
        "template": "esempi/rinascimentale.png",
        "repo_id": "davidmoref/sdxl-lora-adapter-renaissance",
        "weight_name": "pytorch_lora_weights.safetensors",
        "prompt": "renaissance oil painting, oil painting face and hair, elegant brush strokes, chiaroscuro lighting, classic masterpiece, perfectly integrated head",
        "negative_prompt": "ugly, deformed, modern, photography, bad anatomy, worst quality, low quality, pixel art, cartoon"
    },
    "Stile 3: Genio Creativo (Pop-Art)": {
        "template": "esempi/popart.png",
        "repo_id": None,
        "weight_name": None,
        "prompt": "vintage pop art, Comic book style pop-art face and hair, vibrant colors, halftone dots, graphic illustration, bold lines",
        "negative_prompt": "ugly, deformed, photorealistic, realistic photography, oil painting, 3d render"
    }
}

# ==============================================================================
# 3. FUNZIONI CORE DELL'APPLICAZIONE
# ==============================================================================

def create_head_mask(image_pil, face_bbox):
    """Crea una maschera che copre l'intera testa (inclusi capelli e barba) del template."""
    x1, y1, x2, y2 = face_bbox
    w = x2 - x1
    h = y2 - y1
    
    cx, cy = x1 + w/2, y1 + h/2
    mask = Image.new("L", image_pil.size, 0)
    draw = ImageDraw.Draw(mask)
    
    # Area ovale molto abbondante per coprire capelli in alto e barba in basso
    mask_x1 = max(0, cx - w*1.1)
    mask_x2 = min(image_pil.width, cx + w*1.1)
    mask_y1 = max(0, cy - h*1.4) # Su per i capelli/cappelli
    mask_y2 = min(image_pil.height, cy + h*0.8) # Giù per la barba
    
    draw.ellipse((mask_x1, mask_y1, mask_x2, mask_y2), fill=255)
    
    mask = mask.filter(ImageFilter.GaussianBlur(15))
    return mask

@spaces.GPU(duration=60)
def generate_avatar(user_image, selected_style):
    if user_image is None:
        raise gr.Error("Per favore, carica un'immagine del volto.")

    # 1. Caricamento e analisi del Template
    style_config = STYLES[selected_style]
    template_path = style_config["template"]
    
    if not os.path.exists(template_path):
        raise gr.Error(f"Errore: il file di template {template_path} non esiste.")
        
    template_pil = Image.open(template_path).convert("RGB")
    cv_template_image = cv2.cvtColor(np.array(template_pil), cv2.COLOR_RGB2BGR)
    
    template_faces = face_app.get(cv_template_image)
    if len(template_faces) == 0:
        raise gr.Error("Impossibile trovare un volto nel template di base.")
    template_face = sorted(template_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]

    # 2. Analisi del volto utente
    cv_user_image = cv2.cvtColor(np.array(user_image), cv2.COLOR_RGB2BGR)
    user_faces = face_app.get(cv_user_image)
    if len(user_faces) == 0:
        raise gr.Error("Nessun volto rilevato nella tua foto. Riprova con un selfie più chiaro.")
    user_face = sorted(user_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]

    # 3. Creazione maschera per rigenerazione della testa
    print("Creazione maschera per svuotare la testa del template...")
    mask_pil = create_head_mask(template_pil, template_face.bbox)
    
    # Ritaglio del selfie (per dare all'IP-Adapter solo la testa e ignorare lo sfondo della stanza)
    x1, y1, x2, y2 = user_face.bbox
    w = x2 - x1
    h = y2 - y1
    cx, cy = x1 + w/2, y1 + h/2
    crop_size = max(w, h) * 2.5
    ip_x1 = max(0, int(cx - crop_size/2))
    ip_y1 = max(0, int(cy - crop_size/1.5))
    ip_x2 = min(user_image.width, int(cx + crop_size/2))
    ip_y2 = min(user_image.height, int(cy + crop_size/2))
    ip_image = user_image.crop((ip_x1, ip_y1, ip_x2, ip_y2))

    # 4. Rigenerazione Semantica (SDXL + IP-Adapter)
    repo_id = style_config["repo_id"]
    if repo_id:
        try:
            print(f"Caricamento LoRA {repo_id}...")
            pipe.load_lora_weights(repo_id, weight_name=style_config["weight_name"])
        except Exception as e:
            print(f"Errore caricamento LoRA: {e}")
            
    pipe.to(device)

    print("Inizio Generazione Semantica della Testa (SDXL + IP-Adapter)...")
    # Usa una strength quasi al 100% per ignorare i pixel della testa originale sotto la maschera
    generated_image = pipe(
        prompt=style_config["prompt"],
        negative_prompt=style_config["negative_prompt"],
        image=template_pil,
        mask_image=mask_pil,
        ip_adapter_image=ip_image,
        strength=0.99, 
        guidance_scale=8.0,
        num_inference_steps=35,
    ).images[0]

    if repo_id:
        try:
            pipe.unload_lora_weights()
        except:
            pass

    # 5. Iniezione dell'Identità pura (FaceSwap)
    print("Iniezione dell'identità tramite Inswapper...")
    cv_generated = cv2.cvtColor(np.array(generated_image), cv2.COLOR_RGB2BGR)
    generated_faces = face_app.get(cv_generated)
    
    if len(generated_faces) > 0:
        # Trova la faccia appena generata dall'IA
        gen_face = sorted(generated_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]
        # Incolla ESATTAMENTE gli occhi, naso e bocca originali dell'utente sopra la testa generata
        cv_final = swapper.get(cv_generated, gen_face, user_face, paste_back=True)
        final_image = Image.fromarray(cv2.cvtColor(cv_final, cv2.COLOR_BGR2RGB))
    else:
        print("Attenzione: SDXL non ha generato un volto rilevabile. Restituisco l'immagine senza Inswapper.")
        final_image = generated_image

    return final_image

# ==============================================================================
# 4. INTERFACCIA GRADIO (FRONTEND)
# ==============================================================================

custom_css = """
.container { max-width: 1000px; margin: auto; padding-top: 2rem; }
.header-title { text-align: center; color: #1e3a8a; font-family: 'Inter', sans-serif; font-weight: 800; font-size: 2.5rem; margin-bottom: 0.5rem; }
.header-subtitle { text-align: center; color: #4b5563; font-size: 1.1rem; margin-bottom: 2rem; }
.generate-btn { background: linear-gradient(90deg, #10b981 0%, #3b82f6 100%) !important; color: white !important; font-weight: bold !important; font-size: 1.2rem !important; border: none !important; transition: transform 0.2s ease; }
.generate-btn:hover { transform: scale(1.02); }
"""

theme = gr.themes.Soft(primary_hue="blue", secondary_hue="emerald").set(
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_600",
)

with gr.Blocks(css=custom_css, theme=theme) as demo:
    with gr.Column(elem_classes="container"):
        gr.HTML("<h1 class='header-title'>🇮🇹 Italian Builders Avatar Generator 🇮🇹</h1>")
        gr.HTML("<p class='header-subtitle'>Trasforma il tuo selfie e unisciti alla community.</p>")

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil", label="1. Carica il tuo Selfie (Ben illuminato)")

            with gr.Column(scale=1):
                style_selector = gr.Radio(
                    choices=list(STYLES.keys()),
                    label="2. Scegli il tuo Stile",
                    value=list(STYLES.keys())[0]
                )
                gr.Markdown("<br>")
                generate_button = gr.Button("Genera il tuo Avatar 🚀", elem_classes="generate-btn", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(label="Il tuo Avatar", interactive=False)

        generate_button.click(
            fn=generate_avatar,
            inputs=[input_image, style_selector],
            outputs=output_image
        )

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