Update generator.py

generator.py

@@ -26,7 +26,16 @@ class Generator:
         
         return depth_map, lineart_map
 
-    def predict(self, input_image, user_prompt=""):
+    def predict(
+        self, 
+        input_image, 
+        user_prompt="",
+        guidance_scale=1.5,
+        num_inference_steps=6,
+        img2img_strength=0.3,
+        depth_strength=0.3,
+        lineart_strength=0.3
+    ):
         # 1. Pre-process Inputs
         print("Processing Input...")
         processed_image = resize_image_to_1mp(input_image)
@@ -56,12 +65,14 @@ class Generator:
         # ControlNet order: [InstantID, Zoe, LineArt]
         if face_emb is not None:
             print("Face detected: Applying InstantID.")
-            controlnet_conditioning_scale = [0.6, 0.3, 0.3] 
+            # Use strengths from UI
+            controlnet_conditioning_scale = [0.6, depth_strength, lineart_strength] 
             control_guidance_end = [0.3, 0.6, 0.6] # Stop InstantID early
             self.mh.pipeline.set_ip_adapter_scale(0.6) # Set IP-Adapter (likeness) strength
         else:
             print("No face detected: Disabling InstantID.")
-            controlnet_conditioning_scale = [0.0, 0.3, 0.3] # Disable InstantID weight
+            # Use strengths from UI, but keep InstantID at 0.0
+            controlnet_conditioning_scale = [0.0, depth_strength, lineart_strength] 
             control_guidance_end = [0.3, 0.6, 0.6] 
             self.mh.pipeline.set_ip_adapter_scale(0.0)
             
@@ -79,18 +90,20 @@ class Generator:
             control_image=[processed_image, depth_map, lineart_map], # ControlNet inputs
             image_embeds=face_emb,  # Face embedding (or dummy)
             
-            strength=0.3, # <-- Img2Img strength
+            # --- Parameters from UI ---
+            strength=img2img_strength,
+            num_inference_steps=num_inference_steps, 
+            guidance_scale=guidance_scale,
+            # --- End Parameters from UI ---
+            
             controlnet_conditioning_scale=controlnet_conditioning_scale,
             control_guidance_end=control_guidance_end,
             
-            # LCM settings
-            num_inference_steps=6, 
-            guidance_scale=1.5, # <-- CFG Scale
-            
             clip_skip=2,
             
-            # --- LoRA Strength ---
-            cross_attention_kwargs={"scale": 1.25}
+            # --- LoRA Strength REMOVED ---
+            # No longer needed, as LoRA is fused into the model weights
+            # cross_attention_kwargs={"scale": 1.25} 
             
         ).images[0]