Update controlnet_facefix.py

controlnet_facefix.py

@@ -1,41 +1,34 @@
-# controlnet_facefix.py - KORRIGIERTE VERSION FÜR HF SPACES
+# controlnet_facefix.py - BASIEREND AUF DEINEM FUNKTIONIERENDEN CODE
 import torch
 from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
 from PIL import Image
 import time
+import cv2
+import numpy as np
 
 print("="*60)
-print("CONTROLNET_FACEFIX START")
+print("FACE-FIX BASIEREND AUF CONTROLNET_MODULE")
 print("="*60)
 
-# WICHTIG: Keine .to("cuda") beim Laden!
+# WICHTIG: Dieselben Modelle wie in controlnet_module.py!
 _components_loaded = False
-_depth_processor = None
 _controlnet_depth = None  
 _controlnet_pose = None
 _pipeline = None
 
 def _initialize_components():
-    """Lade alle Komponenten mit Fehlerbehandlung"""
-    global _components_loaded, _depth_processor, _controlnet_depth, _controlnet_pose
+    """Lade ControlNets genau wie in controlnet_module.py"""
+    global _components_loaded, _controlnet_depth, _controlnet_pose
     
     if _components_loaded:
         return True
         
     try:
-        print("1. Lade Depth Processor...")
-        from controlnet_aux import ZoeDetector
-        _depth_processor = ZoeDetector.from_pretrained("lllyasviel/ControlNet")
-        print("   ✅ Depth Processor OK")
-    except Exception as e:
-        print(f"   ❌ Depth Processor Fehler: {e}")
-        return False
-    
-    try:
-        print("2. Lade ControlNet Depth...")
+        print("1. Lade ControlNet Depth...")
+        # GLEICHES MODELL wie in controlnet_module.py!
         _controlnet_depth = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11f1e_sd15_depth",
-            torch_dtype=torch.float16  # KEIN .to("cuda") hier!
+            "lllyasviel/sd-controlnet-depth",  # ← HIER ÄNDERN!
+            torch_dtype=torch.float16
         )
         print("   ✅ ControlNet Depth OK")
     except Exception as e:
@@ -43,11 +36,11 @@ def _initialize_components():
         return False
     
     try:
-        print("3. Lade ControlNet OpenPose...")
-        # WICHTIG: KEIN subfolder="faceonly" - das gibt es nicht!
+        print("2. Lade ControlNet OpenPose...")
+        # GLEICHES MODELL wie in controlnet_module.py!
         _controlnet_pose = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_openpose",
-            torch_dtype=torch.float16  # KEIN .to("cuda") hier!
+            "lllyasviel/sd-controlnet-openpose",  # ← HIER ÄNDERN!
+            torch_dtype=torch.float16
         )
         print("   ✅ ControlNet OpenPose OK")
     except Exception as e:
@@ -58,57 +51,63 @@ def _initialize_components():
     print("✅ ALLE KOMPONENTEN GELADEN")
     return True
 
-def _create_control_images(image):
-    """Erstelle Depth und Pose Bilder"""
+def _extract_depth_map(image):
+    """GENAU DIESELBE FUNKTION wie in controlnet_module.py"""
     try:
-        print("   Erstelle Depth Map...")
-        depth_img = _depth_processor(image)
-        
-        print("   Erstelle OpenPose...")
-        from controlnet_aux import OpenposeDetector
-        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
-        pose_img = openpose(image)
-        
-        return pose_img, depth_img
+        img_array = np.array(image.convert("RGB"))
+        gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+        depth_map = cv2.GaussianBlur(gray, (5, 5), 0)
+        depth_rgb = cv2.cvtColor(depth_map, cv2.COLOR_GRAY2RGB)
+        return Image.fromarray(depth_rgb)
     except Exception as e:
-        print(f"   ❌ Control Images Fehler: {e}")
-        return None, None
+        print(f"Depth Map Fehler: {e}")
+        return image.convert("RGB")
+
+def _extract_pose_simple(image):
+    """Einfache Pose-Extraktion basierend auf controlnet_module.py"""
+    try:
+        img_array = np.array(image.convert("RGB"))
+        edges = cv2.Canny(img_array, 100, 200)
+        pose_image = Image.fromarray(edges).convert("RGB")
+        return pose_image
+    except Exception as e:
+        print(f"Pose Extraction Fehler: {e}")
+        return image.convert("RGB")
 
 def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
-    """HAUPTFUNKTION - Optimiert für HF Spaces"""
+    """Face-Fix basierend auf deiner funktionierenden Logik"""
     print("\n" + "🎭"*50)
-    print("FACE-FIX WIRD AUSGEFÜHRT")
+    print("FACE-FIX MIT BEKANNT FUNKTIONIERENDEN MODELLEN")
     print(f"  Model: {model_id}")
     print(f"  Seed: {seed}")
     print("🎭"*50)
     
     start_time = time.time()
     
-    # 1. Komponenten initialisieren
+    # 1. Komponenten initialisieren (mit bekannten Modellen)
     if not _initialize_components():
         print("❌ Komponenten konnten nicht geladen werden")
         return image
     
-    # 2. Control Images erstellen
-    pose_img, depth_img = _create_control_images(image)
-    if pose_img is None or depth_img is None:
-        print("❌ Control Images konnten nicht erstellt werden")
-        return image
+    # 2. Control Images erstellen (mit deinen funktionierenden Methoden)
+    print("🎭 Erstelle Control Images...")
+    depth_img = _extract_depth_map(image)
+    pose_img = _extract_pose_simple(image)
     
-    # 3. Pipeline erstellen (lazy)
+    # 3. Pipeline erstellen
     global _pipeline
     if _pipeline is None:
         try:
-            print("🔄 Lade Pipeline...")
+            print("🔄 Lade Face-Fix Pipeline...")
             _pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 model_id,
-                controlnet=[_controlnet_pose, _controlnet_depth],
+                controlnet=[_controlnet_pose, _controlnet_depth],  # OpenPose zuerst, dann Depth
                 torch_dtype=torch.float16,
                 safety_checker=None,
                 requires_safety_checker=False,
             )
             
-            # HF Spaces Optimierungen
+            # Optimierungen
             _pipeline.enable_attention_slicing()
             _pipeline.enable_vae_slicing()
             
@@ -118,26 +117,26 @@ def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: i
             return image
     
     try:
-        # 4. Auf GPU bewegen (erst JETZT!)
+        # 4. Auf Device bewegen
         device = "cuda" if torch.cuda.is_available() else "cpu"
         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, ugly"
+        face_negative = f"{negative_prompt}, deformed face, blurry face, bad anatomy"
         
-        print("⚡ Führe Face-Fix Inference aus...")
+        print("⚡ Führe Face-Fix aus...")
         
         # 6. Face-Fix ausführen
         result = pipeline(
             prompt=face_prompt,
             negative_prompt=face_negative,
             image=image,
-            mask_image=None,  # Ganzer Bildbereich
+            mask_image=None,
             control_image=[pose_img, depth_img],
             controlnet_conditioning_scale=[0.8, 0.6],  # OpenPose stärker
-            strength=0.4,  # Mittlere Stärke
+            strength=0.4,
             num_inference_steps=20,
             guidance_scale=7.0,
             generator=torch.Generator(device).manual_seed(seed),
@@ -157,6 +156,5 @@ def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: i
         return image
 
 print("="*60)
-print("CONTROLNET_FACEFIX INITIALISIERT")
-print(f"apply_facefix Funktion: {'apply_facefix' in globals()}")
+print("FACE-FIX MODUL FERTIG INITIALISIERT")
 print("="*60)