Update controlnet_module.py

controlnet_module.py

@@ -35,6 +35,7 @@ class ControlNetProcessor:
         self.controlnet_canny = None
         self.pipe_openpose = None
         self.pipe_canny = None
+        self.pipe_multi = None
 
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
@@ -76,9 +77,9 @@ class ControlNetProcessor:
         try:
             img_array = np.array(image.convert("RGB"))
             
-            # Canny Edge Detection
+            # Canny Edge Detection mit besseren Parametern für Gesichter
             gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-            edges = cv2.Canny(gray, 100, 200)
+            edges = cv2.Canny(gray, 50, 150)  # Bessere Kantenerkennung
             
             # Zu 3-Kanal Bild konvertieren
             edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
@@ -142,6 +143,38 @@ class ControlNetProcessor:
                     raise
             return self.pipe_canny
 
+        elif controlnet_type == "multi":
+            if self.pipe_multi is None:
+                print("Loading Multi-ControlNet pipeline...")
+                try:
+                    # Beide ControlNet-Modelle laden
+                    self.controlnet_openpose = ControlNetModel.from_pretrained(
+                        "lllyasviel/sd-controlnet-openpose",
+                        torch_dtype=self.torch_dtype
+                    )
+                    self.controlnet_canny = ControlNetModel.from_pretrained(
+                        "lllyasviel/sd-controlnet-canny", 
+                        torch_dtype=self.torch_dtype
+                    )
+                    
+                    # Multi-ControlNet Pipeline
+                    self.pipe_multi = StableDiffusionControlNetPipeline.from_pretrained(
+                        "runwayml/stable-diffusion-v1-5",
+                        controlnet=[self.controlnet_openpose, self.controlnet_canny],
+                        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 loaded successfully!")
+                except Exception as e:
+                    print(f"Fehler beim Laden von Multi-ControlNet: {e}")
+                    raise
+            return self.pipe_multi
+
     def generate_with_controlnet(
         self, image, prompt, negative_prompt,
         steps, guidance_scale, controlnet_strength,
@@ -149,28 +182,30 @@ class ControlNetProcessor:
     ):
         """Generiert Bild mit ControlNet und Fortschrittsanzeige"""
         try:
-            # --- KORREKTE LOGIK ---
+            # --- MULTI-CONTROLNET LOGIK ---
             if keep_environment:
-                # UMGEBUNG BEIBEHALTEN, PERSON ÄNDERN → KOMBINIERTE STRATEGIE
-                print("🎯 ControlNet Modus: Umgebung beibehalten (OpenPose + Canny Kombination)")
+                # UMGEBUNG BEIBEHALTEN, PERSON ÄNDERN → MULTI-CONTROLNET
+                print("🎯 ControlNet Modus: Umgebung beibehalten (Multi-ControlNet: OpenPose + Canny)")
                 
-                # Schritt 1: OpenPose für Grundpose
+                # Beide Conditioning Maps erstellen
                 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")
+                print("✅ OpenPose + Canny Maps erstellt")
                 
-                # Kombinierte Conditioning - zuerst mit OpenPose arbeiten
-                conditioning_image = pose_image
-                controlnet_type = "openpose"
+                # Multi-ControlNet verwenden
+                conditioning_images = [pose_image, canny_image]
+                controlnet_type = "multi"
+                
+                # Unterschiedliche Strengths für Pose und Canny
+                controlnet_conditioning_scale = [controlnet_strength * 0.7,  # OpenPose: 70%
+                                               controlnet_strength * 0.3]   # Canny: 30%
                 
             else:
-                # PERSON BEIBEHALTEN, UMGEBUNG ÄNDERN → NUR OPENPOSE (wie bisher)
+                # PERSON BEIBEHALTEN, UMGEBUNG ÄNDERN → NUR OPENPOSE
                 controlnet_type = "openpose"
                 print("🎯 ControlNet Modus: Person beibehalten (OpenPose)")
-                conditioning_image = self.extract_pose(image)
+                conditioning_images = self.extract_pose(image)
+                controlnet_conditioning_scale = controlnet_strength
 
             pipe = self.load_controlnet_pipeline(controlnet_type)
 
@@ -185,25 +220,40 @@ class ControlNetProcessor:
             print("🔄 ControlNet: Starte Pipeline...")
 
             # ControlNet Generierung
-            result = pipe(
-                prompt=prompt,
-                image=conditioning_image,
-                negative_prompt=negative_prompt,
-                num_inference_steps=int(steps),
-                guidance_scale=guidance_scale,
-                generator=generator,
-                controlnet_conditioning_scale=controlnet_strength,
-                height=512,
-                width=512,
-                output_type="pil",
-                callback_on_step_end=callback,
-                callback_on_step_end_tensor_inputs=[],
-            )
+            if controlnet_type == "multi":
+                result = pipe(
+                    prompt=prompt,
+                    image=conditioning_images,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=int(steps),
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                    controlnet_conditioning_scale=controlnet_conditioning_scale,
+                    height=512,
+                    width=512,
+                    output_type="pil",
+                    callback_on_step_end=callback,
+                    callback_on_step_end_tensor_inputs=[],
+                )
+            else:
+                result = pipe(
+                    prompt=prompt,
+                    image=conditioning_images,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=int(steps),
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                    controlnet_conditioning_scale=controlnet_conditioning_scale,
+                    height=512,
+                    width=512,
+                    output_type="pil",
+                    callback_on_step_end=callback,
+                    callback_on_step_end_tensor_inputs=[],
+                )
 
             print("✅ ControlNet abgeschlossen!")
             
-            # ZWEI Werte zurückgeben: ControlNet-Output + ORIGINALBILD für Inpaint
-            return result.images[0], image  # ← IMMER Originalbild für Inpaint!
+            return result.images[0], image  # ControlNet-Output + Originalbild
 
         except Exception as e:
             print(f"❌ Fehler in ControlNet: {e}")