Update controlnet_module.py

controlnet_module.py

@@ -1,9 +1,5 @@
 import torch
-from diffusers import (
-    StableDiffusionControlNetPipeline,
-    ControlNetModel,
-    MultiControlNetModel,
-)
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 from controlnet_aux import OpenposeDetector
 from PIL import Image
 import random
@@ -12,9 +8,6 @@ import numpy as np
 import gradio as gr
 
 
-# ============================================================
-# PROGRESS CALLBACK
-# ============================================================
 class ControlNetProgressCallback:
     def __init__(self, progress, total_steps):
         self.progress = progress
@@ -33,36 +26,33 @@ class ControlNetProgressCallback:
         return callback_kwargs
 
 
-# ============================================================
-# CONTROLNET PROZESSOR
-# ============================================================
 class ControlNetProcessor:
     def __init__(self, device="cuda", torch_dtype=torch.float32):
         self.device = device
         self.torch_dtype = torch_dtype
         self.pose_detector = None
-        self.pipe_multi = None  # Multi-ControlNet (OpenPose + Canny)
+        self.controlnet_openpose = None
+        self.controlnet_canny = None
+        self.pipe_openpose = None
+        self.pipe_canny = None
 
-    # ------------------------------------------------------------
-    # POSE DETECTOR
-    # ------------------------------------------------------------
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
         if self.pose_detector is None:
-            print("🧠 Lade Pose Detector...")
+            print("Loading Pose Detector...")
             try:
                 self.pose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
             except Exception as e:
-                print(f"⚠️ Warnung: Pose-Detector konnte nicht geladen werden: {e}")
+                print(f"Warnung: Pose-Detector konnte nicht geladen werden: {e}")
         return self.pose_detector
 
     def extract_pose_simple(self, image):
-        """Fallback: Kantenbasierte Pose"""
+        """Einfache Pose-Extraktion ohne komplexe Abhängigkeiten"""
         try:
             img_array = np.array(image.convert("RGB"))
             edges = cv2.Canny(img_array, 100, 200)
             pose_image = Image.fromarray(edges).convert("RGB")
-            print("⚠️ Verwende einfache Kanten-Pose")
+            print("⚠️ Verwende Kanten-basierte Pose-Approximation")
             return pose_image
         except Exception as e:
             print(f"Fehler bei einfacher Pose-Extraktion: {e}")
@@ -74,99 +64,135 @@ class ControlNetProcessor:
             detector = self.load_pose_detector()
             if detector is None:
                 return self.extract_pose_simple(image)
+
             pose_image = detector(image, hand_and_face=True)
-            print("✅ Pose-Map erfolgreich extrahiert")
             return pose_image
         except Exception as e:
             print(f"Fehler bei Pose-Extraktion: {e}")
             return self.extract_pose_simple(image)
 
-    # ------------------------------------------------------------
-    # CANNY EDGE
-    # ------------------------------------------------------------
     def extract_canny_edges(self, image):
-        """Extrahiert Canny-Kantenbild zur Umgebungserhaltung"""
+        """Extrahiert Canny Edges für Umgebungserhaltung"""
         try:
             img_array = np.array(image.convert("RGB"))
+            
+            # Canny Edge Detection
             gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
             edges = cv2.Canny(gray, 100, 200)
+            
+            # Zu 3-Kanal Bild konvertieren
             edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
             edges_image = Image.fromarray(edges_rgb)
-            print("✅ Canny Edge Map erstellt")
+            
+            print("✅ Canny Edge für Umgebungserhaltung erstellt")
             return edges_image
         except Exception as e:
-            print(f"Fehler bei Canny-Extraktion: {e}")
+            print(f"Fehler bei Canny Edge Extraction: {e}")
             return image.convert("RGB").resize((512, 512))
 
-    # ------------------------------------------------------------
-    # PIPELINE-LADER
-    # ------------------------------------------------------------
-    def load_controlnet_pipeline(self):
-        """Lädt kombinierte Multi-ControlNet Pipeline (OpenPose + Canny)"""
-        if self.pipe_multi is None:
-            print("🧩 Lade Multi-ControlNet Pipeline (OpenPose + Canny)...")
-            try:
-                controlnet_openpose = ControlNetModel.from_pretrained(
-                    "lllyasviel/sd-controlnet-openpose",
-                    torch_dtype=self.torch_dtype
-                )
-                controlnet_canny = ControlNetModel.from_pretrained(
-                    "lllyasviel/sd-controlnet-canny",
-                    torch_dtype=self.torch_dtype
-                )
-
-                multi_controlnet = MultiControlNetModel([controlnet_openpose, controlnet_canny])
-
-                self.pipe_multi = StableDiffusionControlNetPipeline.from_pretrained(
-                    "runwayml/stable-diffusion-v1-5",
-                    controlnet=multi_controlnet,
-                    torch_dtype=self.torch_dtype,
-                    safety_checker=None,
-                    requires_safety_checker=False
-                ).to(self.device)
-
-                from diffusers import EulerAncestralDiscreteScheduler
-                self.pipe_multi.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_multi.scheduler.config)
-                self.pipe_multi.enable_attention_slicing()
-
-                print("✅ Multi-ControlNet Pipeline erfolgreich geladen!")
-            except Exception as e:
-                print(f"❌ Fehler beim Laden von Multi-ControlNet: {e}")
-                raise
-        return self.pipe_multi
+    def load_controlnet_pipeline(self, controlnet_type="openpose"):
+        """Lädt die passende ControlNet Pipeline"""
+        if controlnet_type == "openpose":
+            if self.pipe_openpose is None:
+                print("Loading OpenPose ControlNet pipeline...")
+                try:
+                    self.controlnet_openpose = ControlNetModel.from_pretrained(
+                        "lllyasviel/sd-controlnet-openpose",
+                        torch_dtype=self.torch_dtype
+                    )
+                    self.pipe_openpose = StableDiffusionControlNetPipeline.from_pretrained(
+                        "runwayml/stable-diffusion-v1-5",
+                        controlnet=self.controlnet_openpose,
+                        torch_dtype=self.torch_dtype,
+                        safety_checker=None,
+                        requires_safety_checker=False
+                    ).to(self.device)
+
+                    from diffusers import EulerAncestralDiscreteScheduler
+                    self.pipe_openpose.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_openpose.scheduler.config)
+                    self.pipe_openpose.enable_attention_slicing()
+                    print("✅ OpenPose ControlNet pipeline loaded successfully!")
+                except Exception as e:
+                    print(f"Fehler beim Laden von OpenPose ControlNet: {e}")
+                    raise
+            return self.pipe_openpose
+
+        elif controlnet_type == "canny":
+            if self.pipe_canny is None:
+                print("Loading Canny ControlNet pipeline...")
+                try:
+                    self.controlnet_canny = ControlNetModel.from_pretrained(
+                        "lllyasviel/sd-controlnet-canny",
+                        torch_dtype=self.torch_dtype
+                    )
+                    self.pipe_canny = StableDiffusionControlNetPipeline.from_pretrained(
+                        "runwayml/stable-diffusion-v1-5",
+                        controlnet=self.controlnet_canny,
+                        torch_dtype=self.torch_dtype,
+                        safety_checker=None,
+                        requires_safety_checker=False
+                    ).to(self.device)
+
+                    from diffusers import EulerAncestralDiscreteScheduler
+                    self.pipe_canny.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_canny.scheduler.config)
+                    self.pipe_canny.enable_attention_slicing()
+                    print("✅ Canny ControlNet pipeline loaded successfully!")
+                except Exception as e:
+                    print(f"Fehler beim Laden von Canny ControlNet: {e}")
+                    raise
+            return self.pipe_canny
 
-    # ------------------------------------------------------------
-    # GENERIERUNG
-    # ------------------------------------------------------------
     def generate_with_controlnet(
         self, image, prompt, negative_prompt,
         steps, guidance_scale, controlnet_strength,
-        progress=None
+        progress=None, keep_environment=False
     ):
-        """Generiert Bild mit OpenPose + Canny Kombination"""
+        """Generiert Bild mit ControlNet und Fortschrittsanzeige"""
         try:
-            print("🎯 Modus: Kombiniert (OpenPose + Canny + Inpaint)")
-
-            pose_map = self.extract_pose(image)
-            canny_map = self.extract_canny_edges(image)
-            pipe = self.load_controlnet_pipeline()
-
+            # --- KORREKTE LOGIK ---
+            if keep_environment:
+                # UMGEBUNG BEIBEHALTEN, PERSON ÄNDERN → KOMBINIERTE STRATEGIE
+                print("🎯 ControlNet Modus: Umgebung beibehalten (OpenPose + Canny Kombination)")
+                
+                # Schritt 1: OpenPose für Grundpose
+                pose_image = self.extract_pose(image)
+                print("✅ OpenPose für Grundpose erstellt")
+                
+                # Schritt 2: Canny für Silhouette + Umgebung
+                canny_image = self.extract_canny_edges(image)
+                print("✅ Canny für Silhouette + Umgebung erstellt")
+                
+                # Kombinierte Conditioning - zuerst mit OpenPose arbeiten
+                conditioning_image = pose_image
+                controlnet_type = "openpose"
+                
+            else:
+                # PERSON BEIBEHALTEN, UMGEBUNG ÄNDERN → NUR OPENPOSE (wie bisher)
+                controlnet_type = "openpose"
+                print("🎯 ControlNet Modus: Person beibehalten (OpenPose)")
+                conditioning_image = self.extract_pose(image)
+
+            pipe = self.load_controlnet_pipeline(controlnet_type)
+
+            # Zufälliger Seed
             seed = random.randint(0, 2**32 - 1)
             generator = torch.Generator(device=self.device).manual_seed(seed)
             print(f"ControlNet Seed: {seed}")
 
+            # Fortschritt-Callback
             callback = ControlNetProgressCallback(progress, int(steps)) if progress is not None else None
 
-            print("🚀 Starte kombinierte ControlNet-Pipeline...")
+            print("🔄 ControlNet: Starte Pipeline...")
 
+            # ControlNet Generierung
             result = pipe(
                 prompt=prompt,
-                image=[pose_map, canny_map],  # ← Beide Steuerbilder!
+                image=conditioning_image,
                 negative_prompt=negative_prompt,
                 num_inference_steps=int(steps),
                 guidance_scale=guidance_scale,
-                controlnet_conditioning_scale=[controlnet_strength, controlnet_strength * 0.7],
                 generator=generator,
+                controlnet_conditioning_scale=controlnet_strength,
                 height=512,
                 width=512,
                 output_type="pil",
@@ -174,36 +200,35 @@ class ControlNetProcessor:
                 callback_on_step_end_tensor_inputs=[],
             )
 
-            print("✅ Multi-ControlNet abgeschlossen!")
-            return result.images[0], image  # Für Inpaint weitergeben
+            print("✅ ControlNet abgeschlossen!")
+            
+            # ZWEI Werte zurückgeben: ControlNet-Output + ORIGINALBILD für Inpaint
+            return result.images[0], image  # ← IMMER Originalbild für Inpaint!
 
         except Exception as e:
-            print(f"❌ Fehler in Multi-ControlNet: {e}")
+            print(f"❌ Fehler in ControlNet: {e}")
             import traceback
             traceback.print_exc()
             error_image = image.convert("RGB").resize((512, 512))
             return error_image, error_image
 
-    # ------------------------------------------------------------
-    # INPAINT-VORBEREITUNG
-    # ------------------------------------------------------------
     def prepare_inpaint_input(self, image, keep_environment=False):
         """
         Bereitet das Input-Bild für Inpaint vor
         Rückgabe: (image_für_inpaint, conditioning_info)
         """
         if keep_environment:
+            # PERSON ÄNDERN: Originalbild an Inpaint übergeben
             print("🎯 Inpaint: Übergebe Originalbild (Person ändern)")
             return image, {"type": "original", "image": image}
         else:
+            # UMGEBUNG ÄNDERN: Pose-Map an Inpaint übergeben  
             print("🎯 Inpaint: Übergebe Pose-Map (Umgebung ändern)")
             pose_image = self.extract_pose(image)
             return pose_image, {"type": "pose", "image": pose_image}
 
 
-# ============================================================
-# GLOBALE INSTANZ
-# ============================================================
+# Globale Instanz
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 controlnet_processor = ControlNetProcessor(device=device, torch_dtype=torch_dtype)