Update controlnet_module.py

controlnet_module.py

@@ -160,9 +160,21 @@ class ControlNetProcessor:
             print(f"🔍 Reshaped_input_sizes parameter (in inputs): {inputs.get('reshaped_input_sizes', 'NOT FOUND')}")
 
 
-            # --- NEUEN CODE EINFÜGEN ---
-            # 1. Wähle nur die erste Maske aus den 3 Vorschlägen
-            single_mask = outputs.pred_masks[:, :, 0, :, :]  # Shape wird [1, 1, 256, 256]
+
+            # Debug: Vollständige Dimensionsanalyse
+            print(f"🔍 outputs.pred_masks shape: {outputs.pred_masks.shape}")
+            print(f"🔍 outputs.pred_masks dimensions: {outputs.pred_masks.dim()}")
+
+            # Nach der Auswahl
+            single_mask = outputs.pred_masks[:, :, 0, :, :]
+            print(f"🔍 single_mask shape: {single_mask.shape}")
+            print(f"🔍 single_mask dimensions: {single_mask.dim()}")
+
+            # Wichtig: Ist es wirklich [1, 1, 256, 256] oder [1, 256, 256]?
+            if single_mask.dim() == 3:
+                print("⚠️  WARNUNG: Maske hat nur 3 Dimensionen! Korrigiere...")
+                single_mask = single_mask.unsqueeze(1)  # Fügt Channel-Dimension hinzu
+
 
             # 2. Konvertiere die Größen-Parameter zu CPU und als Python-Liste
             original_size = inputs['original_sizes'].cpu().numpy().tolist()[0]  # [512, 512]