Update controlnet_module.py

controlnet_module.py

@@ -25,7 +25,7 @@ class ControlNetProgressCallback:
         print(f"ControlNet Fortschritt: {self.current_step}/{self.total_steps} ({progress_percentage:.1%})")
         return callback_kwargs
 
-#Initialisierung
+
 class ControlNetProcessor:
     def __init__(self, device="cuda", torch_dtype=torch.float32):
         self.device = device
@@ -39,9 +39,6 @@ class ControlNetProcessor:
         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.conditioning_maps = None
-        self.controlnet_type = None
-        self.controlnet_scales = None
 
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
@@ -257,75 +254,96 @@ class ControlNetProcessor:
                     raise
             return self.pipe_multi_outside
 
+    def generate_with_controlnet(
+        self, image, prompt, negative_prompt,
+        steps, guidance_scale, controlnet_strength,
+        progress=None, keep_environment=False
+    ):
+        """
+        GENERIERT BILD MIT CONTROLNET
+        WICHTIG: Diese Funktion wird von app.py aufgerufen
+        
+        Parameter keep_environment bestimmt:
+        - True: "Umgebung ändern" und "Ausschließlich Gesicht" → Depth+Canny
+        - False: "Focus verändern" → OpenPose+Canny
+        
+        Die eigentliche Maskenlogik wird in app.py (create_face_mask) gehandhabt
+        """
+        try:
+            # --- LOGIK FÜR 3 MODI (VON APP.PY GESTEUERT) ---
+            if keep_environment:
+                # FALL 1 & 3: Umgebung ändern ODER Ausschließlich Gesicht → Depth + Canny
+                print("🎯 ControlNet: Depth + Canny (keep_environment=True)")
+                
+                # Beide Conditioning Maps erstellen
+                depth_image = self.extract_depth_map(image)
+                canny_image = self.extract_canny_edges(image)
+                print("✅ Depth + Canny Maps für Outside/Inside-Box erstellt")
+                
+                # Multi-ControlNet für Outside verwenden
+                conditioning_images = [depth_image, canny_image]
+                controlnet_type = "multi_outside"
+                
+                # 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:
+                # FALL 2: Focus verändern → OpenPose + Canny
+                print("🎯 ControlNet: OpenPose + Canny (keep_environment=False)")
+                
+                # 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}")
+
+            pipe = self.load_controlnet_pipeline(controlnet_type)
+
+            # Fortschritt-Callback
+            callback = ControlNetProgressCallback(progress, int(steps)) if progress is not None else None
 
-    def prepare_conditioning_maps(self, image, keep_environment=False):
-    """
-    NEUE METHODE: Erstellt und speichert Conditioning-Maps basierend auf Bild und Modus.
-    Wird von app.py EINMAL am Anfang aufgerufen.
-    """
-    print("🎯 ControlNet: Erstelle und speichere Conditioning-Maps...")
-    
-    if keep_environment:
-        # FALL: Umgebung ändern oder nur Gesicht -> Depth + Canny
-        print("   Modus: Depth + Canny")
-        self.conditioning_maps = [
-            self.extract_depth_map(image),
-            self.extract_canny_edges(image)
-        ]
-        self.controlnet_type = "multi_outside"
-        self.controlnet_scales = [0.6, 0.4]  # Gewichtung Depth vs Canny
-    else:
-        # FALL: Focus verändern -> OpenPose + Canny
-        print("   Modus: OpenPose + Canny")
-        self.conditioning_maps = [
-            self.extract_pose(image),
-            self.extract_canny_edges(image)
-        ]
-        self.controlnet_type = "multi_inside"
-        self.controlnet_scales = [0.7, 0.3]  # Gewichtung Pose vs Canny
-    
-    print(f"✅ {len(self.conditioning_maps)} Conditioning-Maps gespeichert.")
-    return self.conditioning_maps
-
-def get_controlnet_conditioning(self, noisy_latents, timestep, prompt_embeds, controlnet_strength=1.0):
-    """
-    NEUE KERNMETHODE: Berechnet Steuersignale für einen spezifischen Denoising-Step.
-    Wird von app.py in JEDEM Denoising-Schritt aufgerufen.
-    
-    Args:
-        noisy_latents: Aktuelle verrauschte Latents (Shape: [1, 4, 64, 64])
-        timestep: Aktueller Timestep (z.B. tensor([818]))
-        prompt_embeds: Embeddings des Prompts
-        controlnet_strength: Globale Stärke der ControlNet-Wirkung
-    
-    Returns:
-        controlnet_outputs: Steuersignale, die an UNet übergeben werden
-    """
-    if not hasattr(self, 'conditioning_maps') or self.conditioning_maps is None:
-        raise ValueError("Conditioning-Maps nicht vorhanden. Rufen Sie zuerst prepare_conditioning_maps() auf.")
-    
-    # Lade die passende Multi-ControlNet Pipeline
-    pipe = self.load_controlnet_pipeline(self.controlnet_type)
-    
-    # Skaliere die Conditioning-Stärken mit der globalen controlnet_strength
-    scaled_conditioning_scales = [scale * controlnet_strength for scale in self.controlnet_scales]
-    
-    # WICHTIG: ControlNet-Aufruf für EINEN Step
-    # Die Pipeline muss im "Single-Step-Modus" konfiguriert sein
-    with torch.no_grad():
-        controlnet_outputs = pipe.controlnet(
-            noisy_latents,
-            timestep,
-            encoder_hidden_states=prompt_embeds,
-            controlnet_cond=self.conditioning_maps,
-            conditioning_scale=scaled_conditioning_scales,
-            return_dict=True,
-        )
-    
-    # Gib die Steuersignale zurück (z.B. mid_block_res_sample, down_block_res_samples)
-    return controlnet_outputs
+            print("🔄 ControlNet: Starte Pipeline...")
 
+            # ControlNet Generierung
+            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!")
+            
+            # Rückgabe: ControlNet-Output + Originalbild (für Inpaint)
+            return result.images[0], image
+
+        except Exception as e:
+            print(f"❌ Fehler in ControlNet: {e}")
+            import traceback
+            traceback.print_exc()
+            error_image = image.convert("RGB").resize((512, 512))
+            return error_image, error_image
 
     def prepare_inpaint_input(self, image, keep_environment=False):
         """