Update controlnet_module.py

controlnet_module.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel  # <- KORREKT!
 from controlnet_aux import OpenposeDetector
 from PIL import Image
 import random
@@ -8,8 +8,6 @@ import numpy as np
 import gradio as gr
 
 
-
-
 class ControlNetProgressCallback:
     def __init__(self, progress, total_steps):
         self.progress = progress
@@ -20,7 +18,6 @@ class ControlNetProgressCallback:
         self.current_step = step_index + 1
         progress_percentage = self.current_step / self.total_steps
 
-        # Fortschritt aktualisieren
         if self.progress is not None:
             self.progress(progress_percentage, desc=f"ControlNet: Schritt {self.current_step}/{self.total_steps}")
 
@@ -33,24 +30,52 @@ class ControlNetProcessor:
         self.device = device
         self.torch_dtype = torch_dtype
         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_depth = None
-        self.pipe_multi_inside = None  # OpenPose + Canny für Inside-Box
-        self.pipe_multi_outside = None  # Depth + Canny für Outside-Box
-
+        self.midas_model = None
+        self.midas_transform = None
+        
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
         if self.pose_detector is None:
             print("Loading Pose Detector...")
             try:
                 self.pose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+                print("✅ Pose-Detector geladen")
             except Exception as e:
-                print(f"Warnung: Pose-Detector konnte nicht geladen werden: {e}")
+                print(f"⚠️ Pose-Detector konnte nicht geladen werden: {e}")
         return self.pose_detector
+    
+    def load_midas_model(self):
+        """Lädt MiDaS Model für Depth Maps"""
+        if self.midas_model is None:
+            print("🔄 Lade MiDaS Modell für Depth Maps...")
+            try:
+                # WICHTIG: torchvision 0.20.0 hat MiDaS integriert
+                import torchvision.transforms as T
+                
+                # MiDaS Small (weniger VRAM)
+                self.midas_model = torch.hub.load(
+                    "intel-isl/MiDaS", 
+                    "DPT_Hybrid", 
+                    trust_repo=True
+                )
+                
+                self.midas_model.to(self.device)
+                self.midas_model.eval()
+                
+                # Transform für MiDaS
+                self.midas_transform = T.Compose([
+                    T.Resize(384),
+                    T.ToTensor(),
+                    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+                ])
+                
+                print("✅ MiDaS Modell erfolgreich geladen")
+            except Exception as e:
+                print(f"❌ MiDaS konnte nicht geladen werden: {e}")
+                print("ℹ️ Verwende Fallback-Methode")
+                self.midas_model = None
+        
+        return self.midas_model
 
     def extract_pose_simple(self, image):
         """Einfache Pose-Extraktion ohne komplexe Abhängigkeiten"""
@@ -96,76 +121,67 @@ class ControlNetProcessor:
             print(f"Fehler bei Canny Edge Extraction: {e}")
             return image.convert("RGB").resize((512, 512))
 
-
     def extract_depth_map(self, image):
-    """
-    Extrahiert Depth Map mit MiDaS Small (Fallback auf alten Filter).
-    """
-    try:
-        print("🔄 Versuche MiDaS Small für Depth Map...")
-        # 1. MiDaS Modelle vor dem ersten Gebrauch laden (spart VRAM)
-        if not hasattr(self, 'midas_model'):
-            from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-            import midas
-            
-            self.midas_transform = Compose([
-                Resize(384, interpolation=midas.utils.interpolation),
-                ToTensor(),
-                Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            ])
-            
-            # WICHTIG: MiDaS Small lädt automatisch die 'small'-Variante (weniger VRAM)
-            self.midas_model = midas.MiDaS()
-            self.midas_model.eval()
+        """
+        Extrahiert Depth Map mit MiDaS (Fallback auf Filter)
+        """
+        try:
+            # Versuche MiDaS
+            midas = self.load_midas_model()
+            if midas is not None:
+                print("🎯 Verwende MiDaS für Depth Map...")
+                
+                import torchvision.transforms as T
+                from PIL import Image
+                
+                # Bild vorbereiten
+                img_transformed = self.midas_transform(image).unsqueeze(0).to(self.device)
+                
+                # Depth Map berechnen
+                with torch.no_grad():
+                    prediction = midas(img_transformed)
+                    prediction = torch.nn.functional.interpolate(
+                        prediction.unsqueeze(1),
+                        size=image.size[::-1],  # (height, width)
+                        mode="bicubic",
+                        align_corners=False,
+                    ).squeeze()
+                
+                # Normalisieren für Ausgabe
+                depth_np = prediction.cpu().numpy()
+                depth_min, depth_max = depth_np.min(), depth_np.max()
+                
+                if depth_max > depth_min:
+                    depth_np = (depth_np - depth_min) / (depth_max - depth_min)
+                
+                depth_np = (depth_np * 255).astype(np.uint8)
+                depth_image = Image.fromarray(depth_np).convert("RGB")
+                
+                print("✅ MiDaS Depth Map erfolgreich erstellt")
+                return depth_image
             
-            if self.device == 'cuda':
-                self.midas_model.to(self.device)
-                print("✅ MiDaS Small Modell geladen (GPU)")
             else:
-                print("✅ MiDaS Small Modell geladen (CPU)")
-        
-        # 2. Bild für MiDaS vorbereiten
-        img_input = self.midas_transform(image).unsqueeze(0).to(self.device)
-        
-        # 3. Depth Map berechnen
-        with torch.no_grad():
-            prediction = self.midas_model(img_input)
-            prediction = torch.nn.functional.interpolate(
-                prediction.unsqueeze(1),
-                size=image.size[::-1],  # (height, width)
-                mode="bicubic",
-                align_corners=False,
-            ).squeeze()
-        
-        # 4. Normalisierung für sichtbare Ausgabe
-        depth_np = prediction.cpu().numpy()
-        depth_min, depth_max = depth_np.min(), depth_np.max()
-        if depth_max > depth_min:
-            depth_np = (depth_np - depth_min) / (depth_max - depth_min)
-        depth_np = (depth_np * 255).astype(np.uint8)
-        
-        depth_image = Image.fromarray(depth_np).convert("RGB")
-        print("✅ MiDaS Depth Map erfolgreich erstellt")
-        return depth_image
-        
-    except Exception as e:
-        print(f"⚠️ MiDaS Fehler: {e}. Verwende Fallback (Grayscale Filter)...")
-        # Fallback auf Ihren bestehenden Filter-Code
-        try:
-            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("✅ Fallback Depth Map erstellt")
-            return depth_image
-        except Exception as fallback_error:
-            print(f"❌ Auch Fallback fehlgeschlagen: {fallback_error}")
-            return image.convert("RGB").resize((512, 512))
-
-
+                # Fallback auf einfache Methode
+                print("⚠️ MiDaS nicht verfügbar, verwende Fallback...")
+                raise Exception("MiDaS nicht geladen")
+                
+        except Exception as e:
+            print(f"⚠️ MiDaS Fehler: {e}. Verwende Fallback...")
+            # Fallback auf einfache Depth Map
+            try:
+                img_array = np.array(image.convert("RGB"))
+                gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+
+                # Depth-ähnliche Map erstellen
+                depth_map = cv2.GaussianBlur(gray, (5, 5), 0)
+                depth_rgb = cv2.cvtColor(depth_map, cv2.COLOR_GRAY2RGB)
+                depth_image = Image.fromarray(depth_rgb)
+                
+                print("✅ Fallback Depth Map erstellt")
+                return depth_image
+            except Exception as fallback_error:
+                print(f"❌ Auch Fallback fehlgeschlagen: {fallback_error}")
+                return image.convert("RGB").resize((512, 512))
 
     def prepare_controlnet_maps(self, image, keep_environment=False):
         """
@@ -195,22 +211,15 @@ class ControlNetProcessor:
     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)
-        
-        HINWEIS: Diese Funktion wird nicht direkt von app.py verwendet,
-        da die Logik in generate_with_controlnet enthalten ist.
         """
         if keep_environment:
-            # OUTSIDE-BOX ÄNDERN: Depth+Canny Info für Umgebung
-            print("🎯 Inpaint: Übergebe Depth+Canny Info (Outside-Box ändern)")
+            print("🎯 Inpaint: 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:
-            # INSIDE-BOX ÄNDERN: Originalbild an Inpaint übergeben
-            print("🎯 Inpaint: Übergebe Originalbild (Inside-Box ändern)")
+            print("🎯 Inpaint: Originalbild (Inside-Box ändern)")
             return image, {"type": "original", "image": image}