Update controlnet_module.py

controlnet_module.py

@@ -33,9 +33,12 @@ class ControlNetProcessor:
         self.pose_detector = None
         self.controlnet_openpose = None
         self.controlnet_canny = None
+        self.controlnet_depth = None
         self.pipe_openpose = None
         self.pipe_canny = None
-        self.pipe_multi = None
+        self.pipe_depth = None
+        self.pipe_multi_inside = None  # OpenPose + Canny für Inside-Box
+        self.pipe_multi_outside = None  # Depth + Canny für Outside-Box
 
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
@@ -85,12 +88,31 @@ class ControlNetProcessor:
             edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
             edges_image = Image.fromarray(edges_rgb)
             
-            print("✅ Canny Edge für Umgebungserhaltung erstellt")
+            print("✅ Canny Edge Map erstellt")
             return edges_image
         except Exception as e:
             print(f"Fehler bei Canny Edge Extraction: {e}")
             return image.convert("RGB").resize((512, 512))
 
+    def extract_depth_map(self, image):
+        """Extrahiert Depth Map für räumliche Konsistenz"""
+        try:
+            # Für echte Depth-Maps würde man ein Depth-Estimation-Modell verwenden
+            # Hier als Fallback: Konvertierung zu Grayscale als Depth-Approximation
+            img_array = np.array(image.convert("RGB"))
+            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+            
+            # Depth-ähnliche Map erstellen (helle Bereiche = nah, dunkle = fern)
+            depth_map = cv2.GaussianBlur(gray, (5, 5), 0)
+            depth_rgb = cv2.cvtColor(depth_map, cv2.COLOR_GRAY2RGB)
+            depth_image = Image.fromarray(depth_rgb)
+            
+            print("✅ Depth Map erstellt (Grayscale Approximation)")
+            return depth_image
+        except Exception as e:
+            print(f"Fehler bei Depth Map 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":
@@ -143,11 +165,35 @@ class ControlNetProcessor:
                     raise
             return self.pipe_canny
 
-        elif controlnet_type == "multi":
-            if self.pipe_multi is None:
-                print("Loading Multi-ControlNet pipeline...")
+        elif controlnet_type == "depth":
+            if self.pipe_depth is None:
+                print("Loading Depth ControlNet pipeline...")
+                try:
+                    self.controlnet_depth = ControlNetModel.from_pretrained(
+                        "lllyasviel/sd-controlnet-depth",
+                        torch_dtype=self.torch_dtype
+                    )
+                    self.pipe_depth = StableDiffusionControlNetPipeline.from_pretrained(
+                        "runwayml/stable-diffusion-v1-5",
+                        controlnet=self.controlnet_depth,
+                        torch_dtype=self.torch_dtype,
+                        safety_checker=None,
+                        requires_safety_checker=False
+                    ).to(self.device)
+
+                    from diffusers import EulerAncestralDiscreteScheduler
+                    self.pipe_depth.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_depth.scheduler.config)
+                    self.pipe_depth.enable_attention_slicing()
+                    print("✅ Depth ControlNet pipeline loaded successfully!")
+                except Exception as e:
+                    print(f"Fehler beim Laden von Depth ControlNet: {e}")
+                    raise
+            return self.pipe_depth
+
+        elif controlnet_type == "multi_inside":  # OpenPose + Canny für Inside-Box
+            if self.pipe_multi_inside is None:
+                print("Loading Multi-ControlNet pipeline für Inside-Box...")
                 try:
-                    # Beide ControlNet-Modelle laden
                     if self.controlnet_openpose is None:
                         self.controlnet_openpose = ControlNetModel.from_pretrained(
                             "lllyasviel/sd-controlnet-openpose",
@@ -159,8 +205,7 @@ class ControlNetProcessor:
                             torch_dtype=self.torch_dtype
                         )
                     
-                    # Multi-ControlNet Pipeline
-                    self.pipe_multi = StableDiffusionControlNetPipeline.from_pretrained(
+                    self.pipe_multi_inside = StableDiffusionControlNetPipeline.from_pretrained(
                         "runwayml/stable-diffusion-v1-5",
                         controlnet=[self.controlnet_openpose, self.controlnet_canny],
                         torch_dtype=self.torch_dtype,
@@ -169,13 +214,45 @@ class ControlNetProcessor:
                     ).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!")
+                    self.pipe_multi_inside.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_multi_inside.scheduler.config)
+                    self.pipe_multi_inside.enable_attention_slicing()
+                    print("✅ Multi-ControlNet (Inside) pipeline loaded successfully!")
                 except Exception as e:
-                    print(f"Fehler beim Laden von Multi-ControlNet: {e}")
+                    print(f"Fehler beim Laden von Multi-ControlNet Inside: {e}")
                     raise
-            return self.pipe_multi
+            return self.pipe_multi_inside
+
+        elif controlnet_type == "multi_outside":  # Depth + Canny für Outside-Box
+            if self.pipe_multi_outside is None:
+                print("Loading Multi-ControlNet pipeline für Outside-Box...")
+                try:
+                    if self.controlnet_depth is None:
+                        self.controlnet_depth = ControlNetModel.from_pretrained(
+                            "lllyasviel/sd-controlnet-depth",
+                            torch_dtype=self.torch_dtype
+                        )
+                    if self.controlnet_canny is None:
+                        self.controlnet_canny = ControlNetModel.from_pretrained(
+                            "lllyasviel/sd-controlnet-canny", 
+                            torch_dtype=self.torch_dtype
+                        )
+                    
+                    self.pipe_multi_outside = StableDiffusionControlNetPipeline.from_pretrained(
+                        "runwayml/stable-diffusion-v1-5",
+                        controlnet=[self.controlnet_depth, 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_outside.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe_multi_outside.scheduler.config)
+                    self.pipe_multi_outside.enable_attention_slicing()
+                    print("✅ Multi-ControlNet (Outside) pipeline loaded successfully!")
+                except Exception as e:
+                    print(f"Fehler beim Laden von Multi-ControlNet Outside: {e}")
+                    raise
+            return self.pipe_multi_outside
 
     def generate_with_controlnet(
         self, image, prompt, negative_prompt,
@@ -184,40 +261,45 @@ class ControlNetProcessor:
     ):
         """Generiert Bild mit ControlNet und Fortschrittsanzeige"""
         try:
-            # --- KORRIGIERTE LOGIK ---
+            # --- KORRIGIERTE LOGIK FÜR KONSISTENZ MIT APP.PY ---
             if keep_environment:
-                # UMGEBUNG BEIBEHALTEN, PERSON ÄNDERN → MULTI-CONTROLNET
-                print("🎯 ControlNet Modus: Umgebung beibehalten (Multi-ControlNet: OpenPose + Canny)")
+                # OUTSIDE-BOX ÄNDERN (Umgebung bleibt erhalten) → Depth + Canny
+                print("🎯 ControlNet Modus: Outside-Box ändern (Depth + Canny)")
                 
                 # Beide Conditioning Maps erstellen
-                pose_image = self.extract_pose(image)
+                depth_image = self.extract_depth_map(image)
                 canny_image = self.extract_canny_edges(image)
-                print("✅ OpenPose + Canny Maps erstellt")
-                
-                # Multi-ControlNet verwenden
-                conditioning_images = [pose_image, canny_image]
-                controlnet_type = "multi"
+                print("✅ Depth + Canny Maps für Outside-Box erstellt")
                 
-                # Unterschiedliche Strengths für Pose und Canny
-                controlnet_conditioning_scale = [controlnet_strength * 0.6,  # OpenPose: 60% für Person
-                                               controlnet_strength * 0.4]   # Canny: 40% für Umgebung
+                # Multi-ControlNet für Outside verwenden
+                conditioning_images = [depth_image, canny_image]
+                controlnet_type = "multi_outside"
                 
-                # Zufälliger Seed
-                seed = random.randint(0, 2**32 - 1)
-                generator = torch.Generator(device=self.device).manual_seed(seed)
-                print(f"ControlNet Seed: {seed}")
+                # Gewichtung: Depth 60%, Canny 40%
+                controlnet_conditioning_scale = [controlnet_strength * 0.6,  # Depth: 60% für räumliche Tiefe
+                                               controlnet_strength * 0.4]   # Canny: 40% für Strukturen
                 
             else:
-                # PERSON BEIBEHALTEN, UMGEBUNG ÄNDERN → NUR OPENPOSE
-                controlnet_type = "openpose"
-                print("🎯 ControlNet Modus: Person beibehalten (OpenPose)")
-                conditioning_images = self.extract_pose(image)
-                controlnet_conditioning_scale = controlnet_strength
+                # INSIDE-BOX ÄNDERN (Person bleibt erhalten) → OpenPose + Canny
+                print("🎯 ControlNet Modus: Inside-Box ändern (OpenPose + Canny)")
+                
+                # Beide Conditioning Maps erstellen
+                pose_image = self.extract_pose(image)
+                canny_image = self.extract_canny_edges(image)
+                print("✅ OpenPose + Canny Maps für Inside-Box erstellt")
+                
+                # Multi-ControlNet für Inside verwenden
+                conditioning_images = [pose_image, canny_image]
+                controlnet_type = "multi_inside"
+                
+                # Gewichtung: OpenPose 70%, Canny 30%
+                controlnet_conditioning_scale = [controlnet_strength * 0.7,  # OpenPose: 70% für Person
+                                               controlnet_strength * 0.3]   # Canny: 30% für Konturen
 
-                # Zufälliger Seed
-                seed = random.randint(0, 2**32 - 1)
-                generator = torch.Generator(device=self.device).manual_seed(seed)
-                print(f"ControlNet Seed: {seed}")
+            # Zufälliger Seed
+            seed = random.randint(0, 2**32 - 1)
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+            print(f"ControlNet Seed: {seed}")
 
             pipe = self.load_controlnet_pipeline(controlnet_type)
 
@@ -227,36 +309,20 @@ class ControlNetProcessor:
             print("🔄 ControlNet: Starte Pipeline...")
 
             # ControlNet Generierung
-            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=[],
-                )
+            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!")
             
@@ -275,14 +341,17 @@ class ControlNetProcessor:
         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}
+            # OUTSIDE-BOX ÄNDERN: Depth+Canny Info für Umgebung
+            print("🎯 Inpaint: Übergebe Depth+Canny Info (Outside-Box ändern)")
+            depth_image = self.extract_depth_map(image)
+            canny_image = self.extract_canny_edges(image)
+            # Für Inpaint kann eine kombinierte Map verwendet werden
+            combined_map = Image.blend(depth_image.convert("RGB"), canny_image.convert("RGB"), alpha=0.5)
+            return combined_map, {"type": "depth_canny", "image": combined_map}
         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}
+            # INSIDE-BOX ÄNDERN: Originalbild an Inpaint übergeben
+            print("🎯 Inpaint: Übergebe Originalbild (Inside-Box ändern)")
+            return image, {"type": "original", "image": image}
 
 
 # Globale Instanz