Update controlnet_module.py

controlnet_module.py

@@ -160,14 +160,21 @@ class ControlNetProcessor:
             print(f"🔍 Reshaped_input_sizes parameter (in inputs): {inputs.get('reshaped_input_sizes', 'NOT FOUND')}")
 
 
-            # 5. Maske extrahieren und verarbeiten
-            # outputs.pred_masks enthält die Masken-Logits
-            # post_process_masks stellt die Originalgröße wieder her
+            # --- 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]
+
+            # 2. Konvertiere die Größen-Parameter zu CPU und als Python-Liste
+            original_size = inputs['original_sizes'].cpu().numpy().tolist()[0]  # [512, 512]
+            input_size = inputs['reshaped_input_sizes'].cpu().numpy().tolist()[0]  # [512, 512]
+
+            # 3. Korrekter Aufruf mit angepassten Parametern
             mask = self.sam_processor.post_process_masks(
-                outputs.pred_masks,
-                inputs.original_sizes,
-                inputs.reshaped_input_sizes
-            )[0][0]  # [batch_index][mask_index]
+                single_mask,                    # Nur eine Maske
+                original_sizes=[original_size], # Als Python-Liste
+                reshaped_input_sizes=[input_size] # Als Python-Liste
+            )[0][0]  # [batch][mask]
+
 
             # Sigmoid für Wahrscheinlichkeiten, dann Schwellenwert
             mask = mask.sigmoid().cpu().numpy()