Update controlnet_facefix.py

controlnet_facefix.py

@@ -1,4 +1,4 @@
-# controlnet_facefix.py - BASIEREND AUF DEINEM FUNKTIONIERENDEN CODE
+# controlnet_facefix.py - EINFACHE VERSION (GANZES BILD VERBESSERN)
 import torch
 from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
 from PIL import Image
@@ -7,7 +7,7 @@ import cv2
 import numpy as np
 
 print("="*60)
-print("FACE-FIX BASIEREND AUF CONTROLNET_MODULE")
+print("FACE-FIX: GANZES BILD VERBESSERN")
 print("="*60)
 
 # WICHTIG: Dieselben Modelle wie in controlnet_module.py!
@@ -25,9 +25,8 @@ def _initialize_components():
         
     try:
         print("1. Lade ControlNet Depth...")
-        # GLEICHES MODELL wie in controlnet_module.py!
         _controlnet_depth = ControlNetModel.from_pretrained(
-            "lllyasviel/sd-controlnet-depth",  # ← HIER ÄNDERN!
+            "lllyasviel/sd-controlnet-depth",  # ← FUNKTIONIERT BEI DIR!
             torch_dtype=torch.float16
         )
         print("   ✅ ControlNet Depth OK")
@@ -37,9 +36,8 @@ def _initialize_components():
     
     try:
         print("2. Lade ControlNet OpenPose...")
-        # GLEICHES MODELL wie in controlnet_module.py!
         _controlnet_pose = ControlNetModel.from_pretrained(
-            "lllyasviel/sd-controlnet-openpose",  # ← HIER ÄNDERN!
+            "lllyasviel/sd-controlnet-openpose",  # ← FUNKTIONIERT BEI DIR!
             torch_dtype=torch.float16
         )
         print("   ✅ ControlNet OpenPose OK")
@@ -52,7 +50,7 @@ def _initialize_components():
     return True
 
 def _extract_depth_map(image):
-    """GENAU DIESELBE FUNKTION wie in controlnet_module.py"""
+    """Depth Map wie in controlnet_module.py"""
     try:
         img_array = np.array(image.convert("RGB"))
         gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
@@ -64,7 +62,7 @@ def _extract_depth_map(image):
         return image.convert("RGB")
 
 def _extract_pose_simple(image):
-    """Einfache Pose-Extraktion basierend auf controlnet_module.py"""
+    """Einfache Pose-Extraktion"""
     try:
         img_array = np.array(image.convert("RGB"))
         edges = cv2.Canny(img_array, 100, 200)
@@ -75,24 +73,24 @@ def _extract_pose_simple(image):
         return image.convert("RGB")
 
 def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
-    """Face-Fix basierend auf deiner funktionierenden Logik"""
+    """GANZES BILD verbessern mit ControlNets"""
     print("\n" + "🎭"*50)
-    print("FACE-FIX MIT BEKANNT FUNKTIONIERENDEN MODELLEN")
+    print("FACE-FIX: VERBESSERE GANZES BILD")
     print(f"  Model: {model_id}")
     print(f"  Seed: {seed}")
     print("🎭"*50)
     
     start_time = time.time()
     
-    # 1. Komponenten initialisieren (mit bekannten Modellen)
+    # 1. Komponenten initialisieren
     if not _initialize_components():
         print("❌ Komponenten konnten nicht geladen werden")
         return image
     
-    # 2. Control Images erstellen (mit deinen funktionierenden Methoden)
+    # 2. Control Images erstellen
     print("🎭 Erstelle Control Images...")
-    depth_img = _extract_depth_map(image)
-    pose_img = _extract_pose_simple(image)
+    depth_img = _extract_depth_map(image).resize((512, 512))
+    pose_img = _extract_pose_simple(image).resize((512, 512))
     
     # 3. Pipeline erstellen
     global _pipeline
@@ -101,13 +99,13 @@ def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: i
             print("🔄 Lade Face-Fix Pipeline...")
             _pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 model_id,
-                controlnet=[_controlnet_pose, _controlnet_depth],  # OpenPose zuerst, dann Depth
+                controlnet=[_controlnet_pose, _controlnet_depth],
                 torch_dtype=torch.float16,
                 safety_checker=None,
                 requires_safety_checker=False,
             )
             
-            # Optimierungen
+            # Optimierungen für HF Spaces
             _pipeline.enable_attention_slicing()
             _pipeline.enable_vae_slicing()
             
@@ -122,21 +120,21 @@ def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: i
         print(f"  Device: {device}")
         pipeline = _pipeline.to(device)
         
-        # 5. Prompts optimieren
-        face_prompt = f"{prompt}, perfect face, detailed skin, realistic eyes, sharp focus"
-        face_negative = f"{negative_prompt}, deformed face, blurry face, bad anatomy"
+        # 5. Prompts für allgemeine Verbesserung
+        enhanced_prompt = f"{prompt}, high quality, detailed, sharp focus, professional photography"
+        enhanced_negative = f"{negative_prompt}, blurry, low quality, pixelated, artifacts"
         
-        print("⚡ Führe Face-Fix aus...")
+        print("⚡ Verbessere gesamtes Bild...")
         
-        # 6. Face-Fix ausführen
+        # 6. GANZES BILD verbessern (mask_image=None)
         result = pipeline(
-            prompt=face_prompt,
-            negative_prompt=face_negative,
-            image=image,
-            mask_image=None,
+            prompt=enhanced_prompt,
+            negative_prompt=enhanced_negative,
+            image=image.resize((512, 512)),
+            mask_image=None,  # ← WICHTIG: None = ganzes Bild!
             control_image=[pose_img, depth_img],
-            controlnet_conditioning_scale=[0.8, 0.6],  # OpenPose stärker
-            strength=0.4,
+            controlnet_conditioning_scale=[0.7, 0.5],  # Mittel für subtile Verbesserung
+            strength=0.3,  # Niedrig für feine Anpassungen
             num_inference_steps=20,
             guidance_scale=7.0,
             generator=torch.Generator(device).manual_seed(seed),
@@ -144,17 +142,21 @@ def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: i
             width=512,
         ).images[0]
         
+        # Zurück auf Originalgröße
+        if image.size != (512, 512):
+            result = result.resize(image.size)
+        
         duration = time.time() - start_time
-        print(f"\n✅✅✅ FACE-FIX ERFOLGREICH in {duration:.1f}s ✅✅✅")
+        print(f"\n✅✅✅ BILDVERBESSERUNG ERFOLGREICH in {duration:.1f}s ✅✅✅")
         
         return result
         
     except Exception as e:
-        print(f"\n❌❌❌ FACE-FIX FEHLGESCHLAGEN: {e} ❌❌❌")
+        print(f"\n❌❌❌ FEHLER: {e} ❌❌❌")
         import traceback
         traceback.print_exc()
         return image
 
 print("="*60)
-print("FACE-FIX MODUL FERTIG INITIALISIERT")
+print("FACE-FIX MODUL FERTIG (GANZES BILD VERBESSERN)")
 print("="*60)