Upload 2 files

app.py

@@ -9,7 +9,6 @@ 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
@@ -22,13 +21,18 @@ try:
 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 sessione Rembg
-rembg_session = rembg.new_session()
+# 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...")
@@ -42,6 +46,11 @@ pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
     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 
 
 # ==============================================================================
@@ -75,47 +84,33 @@ STYLES = {
 # 3. FUNZIONI CORE DELL'APPLICAZIONE
 # ==============================================================================
 
-def extract_head_crop(image_pil, face_bbox):
-    """Ritaglia l'area della testa, fermandosi alla base del collo per escludere le spalle."""
+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
-    crop_size = max(w, h) * 3
-    
-    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))
+    mask = Image.new("L", image_pil.size, 0)
+    draw = ImageDraw.Draw(mask)
     
-    # Taglio ancora più drastico: fermiamoci poco sotto il mento per escludere a priori qualsiasi spalla
-    crop_y2 = min(image_pil.height, int(y2 + h * 0.35)) 
-    
-    return image_pil.crop((crop_x1, crop_y1, crop_x2, crop_y2))
-
-def cut_bottom(image_rgba, cut_height_ratio=0.20):
-    """Rende completamente trasparente la base dell'immagine (Hard Cut) per eliminare il collo inferiore."""
-    width, height = image_rgba.size
-    cut_height = int(height * cut_height_ratio)
-    if cut_height == 0:
-        return image_rgba
-        
-    r, g, b, a = image_rgba.split()
-    alpha_np = np.array(a)
+    # 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
     
-    # Taglio netto: zero alpha nella parte inferiore
-    alpha_np[-cut_height:, :] = 0
+    draw.ellipse((mask_x1, mask_y1, mask_x2, mask_y2), fill=255)
     
-    a = Image.fromarray(alpha_np)
-    image_rgba.putalpha(a)
-    return image_rgba
+    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 (lo facciamo prima per sapere dove guarda)
+    # 1. Caricamento e analisi del Template
     style_config = STYLES[selected_style]
     template_path = style_config["template"]
     
@@ -130,91 +125,49 @@ def generate_avatar(user_image, selected_style):
         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 iniziale del volto utente per controllare l'orientamento
+    # 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. Controllo Yaw (Orientamento Destra/Sinistra)
-    user_yaw = user_face.pose[1]
-    template_yaw = template_face.pose[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)
     
-    # Se guardano in direzioni opposte, ribaltiamo il selfie
-    if (user_yaw * template_yaw) < 0:
-        print("Orientamenti opposti! Eseguo specchio (FLIP_LEFT_RIGHT) del selfie.")
-        user_image = user_image.transpose(Image.FLIP_LEFT_RIGHT)
-        # Rianalizziamo il volto ribaltato per avere le nuove coordinate perfette
-        cv_user_image = cv2.cvtColor(np.array(user_image), cv2.COLOR_RGB2BGR)
-        user_faces = face_app.get(cv_user_image)
-        user_face = sorted(user_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]
-
-    # 4. Ritaglio, rimozione sfondo e Hard Cut del collo
-    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)
-    head_rgba = cut_bottom(head_rgba, cut_height_ratio=0.20) # Taglio netto trasparente
-    
-    # 5. 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:
-        head_face_crop = user_face
-    else:
-        head_face_crop = sorted(head_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]))[-1]
-        
-    H_bbox = head_face_crop.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
-    
-    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
-
-    # 6. Compositing (Incolla la testa sul template)
-    print("Compositing della testa sul template...")
-    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)
-    
-    # 7. 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))
-    mask_pil = mask_pil.filter(ImageFilter.MaxFilter(7))
-    mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(15))
-    
-    # 8. Stylization tramite Inpainting
+    # 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"Scaricamento/Caricamento LoRA {repo_id}...")
+            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 da HF: {e}")
+            print(f"Errore caricamento LoRA: {e}")
             
     pipe.to(device)
 
-    print("Inizio Inpainting per fondere stile e testa...")
-    stylized_image = pipe(
+    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=composited_pil,
+        image=template_pil,
         mask_image=mask_pil,
-        strength=0.45, 
+        ip_adapter_image=ip_image,
+        strength=0.99, 
         guidance_scale=8.0,
         num_inference_steps=35,
     ).images[0]
@@ -225,7 +178,22 @@ def generate_avatar(user_image, selected_style):
         except:
             pass
 
-    return stylized_image
+    # 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)

requirements.txt

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