Update controlnet_facefix.py

controlnet_facefix.py

@@ -1,67 +1,162 @@
+# controlnet_facefix.py - KORRIGIERTE VERSION FÜR HF SPACES
 import torch
 from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
-from controlnet_aux import ZoeDetector
 from PIL import Image
 import time
 
-print("Lade Face-Fix System (Depth-basiert)...")
+print("="*60)
+print("CONTROLNET_FACEFIX START")
+print("="*60)
 
-# Vereinfachte Version NUR mit Depth
-try:
-    depth_processor = ZoeDetector.from_pretrained("lllyasviel/ControlNet")
-    print("✅ Depth Processor geladen")
+# WICHTIG: Keine .to("cuda") beim Laden!
+_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
+    
+    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...")
+        _controlnet_depth = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11f1e_sd15_depth",
+            torch_dtype=torch.float16  # KEIN .to("cuda") hier!
+        )
+        print("   ✅ ControlNet Depth OK")
+    except Exception as e:
+        print(f"   ❌ ControlNet Depth Fehler: {e}")
+        return False
+    
+    try:
+        print("3. Lade ControlNet OpenPose...")
+        # WICHTIG: KEIN subfolder="faceonly" - das gibt es nicht!
+        _controlnet_pose = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_openpose",
+            torch_dtype=torch.float16  # KEIN .to("cuda") hier!
+        )
+        print("   ✅ ControlNet OpenPose OK")
+    except Exception as e:
+        print(f"   ❌ ControlNet OpenPose Fehler: {e}")
+        return False
+    
+    _components_loaded = True
+    print("✅ ALLE KOMPONENTEN GELADEN")
+    return True
+
+def _create_control_images(image):
+    """Erstelle Depth und Pose Bilder"""
+    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
+    except Exception as e:
+        print(f"   ❌ Control Images Fehler: {e}")
+        return None, None
+
+def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
+    """HAUPTFUNKTION - Optimiert für HF Spaces"""
+    print("\n" + "🎭"*50)
+    print("FACE-FIX WIRD AUSGEFÜHRT")
+    print(f"  Model: {model_id}")
+    print(f"  Seed: {seed}")
+    print("🎭"*50)
+    
+    start_time = time.time()
     
-    controlnet_depth = ControlNetModel.from_pretrained(
-        "lllyasviel/control_v11f1e_sd15_depth",
-        torch_dtype=torch.float16
-    ).to("cuda")
-    print("✅ ControlNet Depth geladen")
+    # 1. Komponenten initialisieren
+    if not _initialize_components():
+        print("❌ Komponenten konnten nicht geladen werden")
+        return image
     
-    _facefix_pipe = None
+    # 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
     
-    def _get_facefix_pipeline(model_id: str):
-        global _facefix_pipe
-        if _facefix_pipe is None:
-            print(f"Lade Face-Fix-Pipeline für Modell: {model_id}")
-            _facefix_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+    # 3. Pipeline erstellen (lazy)
+    global _pipeline
+    if _pipeline is None:
+        try:
+            print("🔄 Lade Pipeline...")
+            _pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 model_id,
-                controlnet=[controlnet_depth],
+                controlnet=[_controlnet_pose, _controlnet_depth],
                 torch_dtype=torch.float16,
                 safety_checker=None,
-            ).to("cuda")
-        return _facefix_pipe
-    
-    def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
-        try:
-            pipe = _get_facefix_pipeline(model_id)
-            depth_img = depth_processor(image)
-            
-            face_prompt = prompt + ", perfect face, detailed skin"
-            face_negative = negative_prompt + ", deformed face"
+                requires_safety_checker=False,
+            )
             
-            fixed = pipe(
-                prompt=face_prompt,
-                negative_prompt=face_negative,
-                image=image,
-                mask_image=None,
-                control_image=[depth_img],
-                controlnet_conditioning_scale=[0.7],
-                strength=0.35,
-                num_inference_steps=20,
-                guidance_scale=7.0,
-                generator=torch.Generator("cuda").manual_seed(seed),
-            ).images[0]
-            
-            return fixed
+            # HF Spaces Optimierungen
+            _pipeline.enable_attention_slicing()
+            _pipeline.enable_vae_slicing()
             
+            print("✅ Pipeline geladen")
         except Exception as e:
-            print(f"Face-Fix Error: {e}")
+            print(f"❌ Pipeline Fehler: {e}")
             return image
-            
-except Exception as e:
-    print(f"❌ Face-Fix Setup fehlgeschlagen: {e}")
     
-    # Fallback: Dummy-Funktion die nichts macht
-    def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
-        print("⚠️ Face-Fix nicht verfügbar, gebe Original zurück")
-        return image
+    try:
+        # 4. Auf GPU bewegen (erst JETZT!)
+        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"
+        
+        print("⚡ Führe Face-Fix Inference aus...")
+        
+        # 6. Face-Fix ausführen
+        result = pipeline(
+            prompt=face_prompt,
+            negative_prompt=face_negative,
+            image=image,
+            mask_image=None,  # Ganzer Bildbereich
+            control_image=[pose_img, depth_img],
+            controlnet_conditioning_scale=[0.8, 0.6],  # OpenPose stärker
+            strength=0.4,  # Mittlere Stärke
+            num_inference_steps=20,
+            guidance_scale=7.0,
+            generator=torch.Generator(device).manual_seed(seed),
+            height=512,
+            width=512,
+        ).images[0]
+        
+        duration = time.time() - start_time
+        print(f"\n✅✅✅ FACE-FIX ERFOLGREICH in {duration:.1f}s ✅✅✅")
+        
+        return result
+        
+    except Exception as e:
+        print(f"\n❌❌❌ FACE-FIX FEHLGESCHLAGEN: {e} ❌❌❌")
+        import traceback
+        traceback.print_exc()
+        return image
+
+print("="*60)
+print("CONTROLNET_FACEFIX INITIALISIERT")
+print(f"apply_facefix Funktion: {'apply_facefix' in globals()}")
+print("="*60)