Update controlnet_module.py

controlnet_module.py

@@ -1,3 +1,4 @@
+# controlnet_processor.py
 import torch
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 from controlnet_aux import OpenposeDetector
@@ -31,8 +32,10 @@ class ControlNetProcessor:
         self.device = device
         self.torch_dtype = torch_dtype
         self.pose_detector = None
-        self.controlnet = None
-        self.pipe = None
+        self.controlnet_openpose = None
+        self.controlnet_canny = None
+        self.pipe_openpose = None
+        self.pipe_canny = None
 
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
@@ -63,12 +66,83 @@ class ControlNetProcessor:
             if detector is None:
                 return self.extract_pose_simple(image)
 
-            pose_image = detector.detect(image, hand_and_face=True)
+            pose_image = detector(image, hand_and_face=True)
             return pose_image
         except Exception as e:
             print(f"Fehler bei Pose-Extraktion: {e}")
             return self.extract_pose_simple(image)
 
+    def extract_canny_edges(self, image):
+        """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 für Umgebungserhaltung erstellt")
+            return edges_image
+        except Exception as e:
+            print(f"Fehler bei Canny Edge Extraction: {e}")
+            return image.convert("RGB").resize((512, 512))
+
+    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
+
     def generate_with_controlnet(
         self, image, prompt, negative_prompt,
         steps, guidance_scale, controlnet_strength,
@@ -76,21 +150,19 @@ class ControlNetProcessor:
     ):
         """Generiert Bild mit ControlNet und Fortschrittsanzeige"""
         try:
-            pipe = self.load_controlnet_pipeline()
-
-            print("🔄 ControlNet: Extrahiere Pose...")
-            if progress:
-                progress(0.05, desc="ControlNet: Extrahiere Pose...")
-
-            # --- Fallunterscheidung ---
+            # --- ENTSCHEIDUNG: Welches ControlNet für welche Aufgabe? ---
             if keep_environment:
-                print("🎯 Modus: Umgebung beibehalten (nutze Originalbild als Quelle)")
-                input_image = image
-                conditioning_image = None
+                # PERSON ÄNDERN, UMGEBUNG BEIBEHALTEN
+                controlnet_type = "canny"
+                print("🎯 ControlNet Modus: Umgebung beibehalten (Canny Edge)")
+                conditioning_image = self.extract_canny_edges(image)
             else:
-                print("🎯 Modus: Umgebung darf sich ändern (nutze Pose-Map)")
+                # UMGEBUNG ÄNDERN, PERSON BEIBEHALTEN  
+                controlnet_type = "openpose"
+                print("🎯 ControlNet Modus: Person beibehalten (OpenPose)")
                 conditioning_image = self.extract_pose(image)
-                input_image = conditioning_image
+
+            pipe = self.load_controlnet_pipeline(controlnet_type)
 
             # Zufälliger Seed
             seed = random.randint(0, 2**32 - 1)
@@ -102,38 +174,21 @@ class ControlNetProcessor:
 
             print("🔄 ControlNet: Starte Pipeline...")
 
-            if conditioning_image is not None:
-                # Umgebung darf sich ändern
-                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=[],
-                )
-            else:
-                # Umgebung soll beibehalten werden
-                result = pipe(
-                    prompt=prompt,
-                    image=input_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=[],
-                )
+            # 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=[],
+            )
 
             # Debug-Ausgabe Scheduler Steps
             try:
@@ -145,41 +200,31 @@ class ControlNetProcessor:
                 print(f"⚠️ Konnte ControlNet Scheduler-Info nicht auslesen: {e}")
 
             print("✅ ControlNet abgeschlossen!")
-            return result.images[0]
+            
+            # KORREKTUR: ZWEI Werte zurückgeben
+            return result.images[0], conditioning_image
 
         except Exception as e:
             print(f"❌ Fehler in ControlNet: {e}")
             import traceback
             traceback.print_exc()
-            return image.convert("RGB").resize((512, 512))
-
-    def load_controlnet_pipeline(self):
-        """Lädt die ControlNet Pipeline"""
-        if self.pipe is None:
-            print("Loading ControlNet pipeline...")
-            try:
-                self.controlnet = ControlNetModel.from_pretrained(
-                    "lllyasviel/sd-controlnet-openpose",
-                    torch_dtype=self.torch_dtype
-                )
-                self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
-                    "runwayml/stable-diffusion-v1-5",
-                    controlnet=self.controlnet,
-                    torch_dtype=self.torch_dtype,
-                    safety_checker=None,
-                    requires_safety_checker=False
-                ).to(self.device)
-
-                # Scheduler wechseln zu Euler Ancestral
-                from diffusers import EulerAncestralDiscreteScheduler
-                self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
-
-                self.pipe.enable_attention_slicing()
-                print("✅ ControlNet pipeline loaded successfully with EulerAncestralDiscreteScheduler!")
-            except Exception as e:
-                print(f"Fehler beim Laden von ControlNet: {e}")
-                raise
-        return self.pipe
+            error_image = image.convert("RGB").resize((512, 512))
+            return error_image, error_image
+
+    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