Update generator.py

generator.py

@@ -20,23 +20,30 @@ class Generator:
         
         # Generate lineart map
         lineart_map_raw = self.mh.lineart_anime_detector(image)
+
+        # Generate color map
+        color_map_raw = self.mh.color_detector(image)
         # --- END MODIFIED ---
         
         # Manually resize maps to match the exact output resolution
         depth_map = depth_map_raw.resize((width, height), Image.LANCZOS)
         lineart_map = lineart_map_raw.resize((width, height), Image.LANCZOS)
+        color_map = color_map_raw.resize((width, height), Image.LANCZOS)
         
-        return depth_map, lineart_map
+        return depth_map, lineart_map, color_map
 
     def predict(
         self, 
         input_image, 
         user_prompt="",
+        negative_prompt="",           # <-- ADDED
         guidance_scale=1.5,
         num_inference_steps=6,
         img2img_strength=0.3,
         depth_strength=0.3,
-        lineart_strength=0.3
+        lineart_strength=0.3,
+        color_strength=0.7,           # <-- ADDED
+        seed=-1                       # <-- ADDED
     ):
         # 1. Pre-process Inputs
         print("Processing Input...")
@@ -58,39 +65,46 @@ class Generator:
             final_prompt = f"{Config.STYLE_TRIGGER}, {user_prompt}"
             
         print(f"Prompt: {final_prompt}")
+        print(f"Negative Prompt: {negative_prompt}")
 
         # 4. Generate Control Maps (Structure)
-        print("Generating Control Maps (Depth, LineArt)...")
-        depth_map, lineart_map = self.prepare_control_images(processed_image, target_width, target_height)
+        print("Generating Control Maps (Depth, LineArt, Color)...")
+        depth_map, lineart_map, color_map = self.prepare_control_images(processed_image, target_width, target_height)
         
         # 5. Logic for Face vs No-Face
-        # ControlNet order: [InstantID, Zoe, LineArt]
+        # --- MODIFIED: Added Color Control ---
+        # ControlNet order: [InstantID, Zoe, LineArt, Color]
         if face_emb is not None:
             print("Face detected: Applying InstantID.")
-            # 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
+            controlnet_conditioning_scale = [0.6, depth_strength, lineart_strength, color_strength] 
+            control_guidance_end = [0.3, 0.6, 0.6, 0.9] # Stop InstantID early, let color run longer
+            self.mh.pipeline.set_ip_adapter_scale(0.6)
         else:
             print("No face detected: Disabling InstantID.")
-            # 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] 
+            controlnet_conditioning_scale = [0.0, depth_strength, lineart_strength, color_strength] 
+            control_guidance_end = [0.3, 0.6, 0.6, 0.9] 
             self.mh.pipeline.set_ip_adapter_scale(0.0)
             
             # --- START FIX for NoneType Error ---
-            # Create a dummy tensor instead of passing None
-            # Shape is (batch_size, embedding_dim)
             face_emb = torch.zeros((1, 512), dtype=Config.DTYPE, device=Config.DEVICE)
             # --- END FIX ---
 
+        # --- ADDED: Seed/Generator Logic ---
+        if seed == -1 or seed is None:
+            seed = torch.Generator().seed()
+        generator = torch.Generator(device=Config.DEVICE).manual_seed(int(seed))
+        print(f"Using seed: {seed}")
+        # --- END ADDED ---
+
         # 6. Run Inference
         print("Running pipeline...")
         result = self.mh.pipeline(
             prompt=final_prompt,
-            image=processed_image,  # Base image for Img2Img
-            control_image=[processed_image, depth_map, lineart_map], # ControlNet inputs
-            image_embeds=face_emb,  # Face embedding (or dummy)
+            negative_prompt=negative_prompt, # <-- ADDED
+            image=processed_image,
+            control_image=[processed_image, depth_map, lineart_map, color_map], # <-- MODIFIED
+            image_embeds=face_emb,
+            generator=generator,             # <-- ADDED
             
             # --- Parameters from UI ---
             strength=img2img_strength,
@@ -103,10 +117,6 @@ class Generator:
             
             clip_skip=2,
             
-            # --- LoRA Strength REMOVED ---
-            # No longer needed, as LoRA is fused into the model weights
-            # cross_attention_kwargs={"scale": 1.25} 
-            
         ).images[0]
         
         return result