Update controlnet_module.py

controlnet_module.py

@@ -94,26 +94,77 @@ 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 für räumliche Konsistenz"""
+    """
+    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()
+            
+            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:
-            # 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-ä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)")
+            print("✅ Fallback Depth Map erstellt")
             return depth_image
-        except Exception as e:
-            print(f"Fehler bei Depth Map Extraction: {e}")
+        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):
         """
         ERSTELLT NUR CONDITIONING-MAPS, generiert KEIN Bild.