Upload 2 files

app.py

@@ -9,6 +9,7 @@ from diffusers import StableDiffusionXLInpaintPipeline
 from insightface.app import FaceAnalysis
 import insightface
 from huggingface_hub import hf_hub_download, snapshot_download
+import rembg
 
 # ==============================================================================
 # 1. SETUP DELL'AMBIENTE E DOWNLOAD MODELLI
@@ -26,14 +27,8 @@ print("Inizializzazione FaceAnalysis...")
 face_app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
 face_app.prepare(ctx_id=0, det_size=(640, 640))
 
-# Download e caricamento del modello Inswapper
-print("Download e caricamento Inswapper...")
-try:
-    inswapper_path = hf_hub_download(repo_id="ezioruan/inswapper_128.onnx", filename="inswapper_128.onnx")
-    face_swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
-except Exception as e:
-    print(f"Errore caricamento Inswapper: {e}")
-    face_swapper = None
+# Inizializza sessione Rembg
+rembg_session = rembg.new_session()
 
 # Caricamento Pipeline Inpainting SDXL
 print("Caricamento pipeline SDXL Inpaint...")
@@ -53,27 +48,26 @@ pipe.watermark = None
 # ==============================================================================
 # 2. CONFIGURAZIONE STILI E TEMPLATE
 # ==============================================================================
-# I template sono le immagini pre-generate fornite dall'utente.
 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, highly detailed, perfect pixel art shading, 8-bit retro aesthetic, matching lighting",
+        "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, elegant brush strokes, chiaroscuro lighting, classic masterpiece, perfectly integrated face",
+        "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, vibrant colors, halftone dots, graphic illustration, bold lines",
+        "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"
     }
 }
@@ -82,37 +76,26 @@ STYLES = {
 # 3. FUNZIONI CORE DELL'APPLICAZIONE
 # ==============================================================================
 
-def create_face_mask(image, bbox, expansion=0.2, blur_radius=15):
-    """
-    Crea una maschera sfumata attorno al volto per l'inpainting.
-    """
-    mask = Image.new("L", image.size, 0)
-    draw = ImageDraw.Draw(mask)
-    
-    x1, y1, x2, y2 = bbox
+def extract_head_crop(image_pil, face_bbox):
+    """Ritaglia un'area abbondante attorno al volto per catturare capelli e collo."""
+    x1, y1, x2, y2 = face_bbox
     w = x2 - x1
     h = y2 - y1
     
-    # Espandi leggermente la bounding box
-    nx1 = max(0, x1 - w * expansion)
-    ny1 = max(0, y1 - h * expansion)
-    nx2 = min(image.width, x2 + w * expansion)
-    ny2 = min(image.height, y2 + h * (expansion * 1.5)) # Espandi un po' di più verso il basso per il mento
+    cx, cy = x1 + w/2, y1 + h/2
+    crop_size = max(w, h) * 3
     
-    # Disegna un ovale
-    draw.ellipse([nx1, ny1, nx2, ny2], fill=255)
+    crop_x1 = max(0, int(cx - crop_size/1.5))
+    crop_y1 = max(0, int(cy - crop_size/1.2)) # Più spazio sopra per i capelli
+    crop_x2 = min(image_pil.width, int(cx + crop_size/1.5))
+    crop_y2 = min(image_pil.height, int(cy + crop_size/2.5)) # Meno spazio sotto per evitare troppe spalle
     
-    # Sfuma i bordi per fondere meglio
-    mask = mask.filter(ImageFilter.GaussianBlur(blur_radius))
-    return mask
+    return image_pil.crop((crop_x1, crop_y1, crop_x2, crop_y2))
 
 @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.")
-        
-    if face_swapper is None:
-        raise gr.Error("Modello Inswapper non caricato correttamente.")
 
     # 1. Analisi del volto utente
     cv_user_image = cv2.cvtColor(np.array(user_image), cv2.COLOR_RGB2BGR)
@@ -120,10 +103,29 @@ def generate_avatar(user_image, selected_style):
     if len(user_faces) == 0:
         raise gr.Error("Nessun volto rilevato nella tua foto. Riprova con un selfie più chiaro.")
     
-    # Prendi il volto più grande
     user_face = sorted(user_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]
 
-    # 2. Caricamento e analisi del Template
+    # 2. Ritaglio e rimozione sfondo (Head Extraction)
+    print("Estrazione della testa con Rembg...")
+    cropped_user = extract_head_crop(user_image, user_face.bbox)
+    head_rgba = rembg.remove(cropped_user, session=rembg_session)
+    
+    # 3. Analisi della testa scontornata per trovare le nuove coordinate del volto
+    head_rgb_cv = cv2.cvtColor(np.array(head_rgba.convert("RGB")), cv2.COLOR_RGB2BGR)
+    head_faces = face_app.get(head_rgb_cv)
+    if len(head_faces) == 0:
+        # Fallback se non trova il volto dopo il ritaglio
+        head_face = user_face
+    else:
+        head_face = sorted(head_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]
+        
+    H_bbox = head_face.bbox
+    H_w = H_bbox[2] - H_bbox[0]
+    H_h = H_bbox[3] - H_bbox[1]
+    H_cx = H_bbox[0] + H_w/2
+    H_cy = H_bbox[1] + H_h/2
+
+    # 4. Caricamento e analisi del Template
     style_config = STYLES[selected_style]
     template_path = style_config["template"]
     
@@ -137,19 +139,38 @@ def generate_avatar(user_image, selected_style):
     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]
-
-    # 3. Swap del Volto (Fotorealistico)
-    print("Eseguendo Face Swap...")
-    swapped_cv = face_swapper.get(cv_template_image, template_face, user_face, paste_back=True)
-    swapped_pil = Image.fromarray(cv2.cvtColor(swapped_cv, cv2.COLOR_BGR2RGB))
     
-    # 4. Creazione Maschera per l'Inpainting
-    mask_pil = create_face_mask(template_pil, template_face.bbox)
+    T_bbox = template_face.bbox
+    T_w = T_bbox[2] - T_bbox[0]
+    T_cx = T_bbox[0] + T_w/2
+    T_cy = T_bbox[1] + (T_bbox[3] - T_bbox[1])/2
+
+    # 5. Compositing (Incolla la testa sul template)
+    print("Compositing della testa sul template...")
+    # Scaliamo la testa in modo che la larghezza del volto coincida (con un piccolo margine in più per coprire bene)
+    scale = (T_w / H_w) * 1.05 
+    new_size = (int(head_rgba.width * scale), int(head_rgba.height * scale))
+    scaled_head = head_rgba.resize(new_size, Image.LANCZOS)
+    
+    new_H_cx = H_cx * scale
+    new_H_cy = H_cy * scale
+    
+    paste_x = int(T_cx - new_H_cx)
+    paste_y = int(T_cy - new_H_cy)
+    
+    composited_pil = template_pil.copy()
+    composited_pil.paste(scaled_head, (paste_x, paste_y), scaled_head)
+    
+    # 6. Creazione della Maschera per l'Inpainting
+    mask_pil = Image.new("L", composited_pil.size, 0)
+    mask_pil.paste(scaled_head.split()[3], (paste_x, paste_y))
+    # Espandi leggermente la maschera per coprire bene i bordi e sfumala
+    mask_pil = mask_pil.filter(ImageFilter.MaxFilter(7))
+    mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(15))
     
-    # 5. Stylization tramite Inpainting
+    # 7. Stylization tramite Inpainting
     repo_id = style_config["repo_id"]
     
-    # Carica il LoRA dello stile se necessario
     if repo_id:
         try:
             print(f"Scaricamento/Caricamento LoRA {repo_id}...")
@@ -157,21 +178,19 @@ def generate_avatar(user_image, selected_style):
         except Exception as e:
             print(f"Errore caricamento LoRA da HF: {e}")
             
-    # Sposta la pipeline su GPU per l'inferenza
     pipe.to(device)
 
-    print("Inizio Inpainting per fondere lo stile...")
+    print("Inizio Inpainting per fondere stile e testa...")
     stylized_image = pipe(
         prompt=style_config["prompt"],
         negative_prompt=style_config["negative_prompt"],
-        image=swapped_pil,
+        image=composited_pil,
         mask_image=mask_pil,
-        strength=0.35, # Bassa strength: mantiene forte l'identità, applica solo lo stile superficiale
-        guidance_scale=7.5,
-        num_inference_steps=30,
+        strength=0.45, # Aumentato leggermente per permettere all'IA di ridisegnare capelli e occhiali nello stile
+        guidance_scale=8.0,
+        num_inference_steps=35,
     ).images[0]
 
-    # Scarica il LoRA per evitare conflitti al prossimo giro
     if repo_id:
         try:
             pipe.unload_lora_weights()
@@ -204,7 +223,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
 
         with gr.Row():
             with gr.Column(scale=1):
-                input_image = gr.Image(type="pil", label="1. Carica il tuo Selfie")
+                input_image = gr.Image(type="pil", label="1. Carica il tuo Selfie (Ben illuminato)")
 
             with gr.Column(scale=1):
                 style_selector = gr.Radio(

requirements.txt

@@ -9,3 +9,4 @@ onnxruntime
 opencv-python
 huggingface-hub
 peft
+rembg