Update generator.py

generator.py

@@ -1,6 +1,6 @@
 import torch
 from config import Config
-from utils import resize_image_to_1mp, get_caption
+from utils import resize_image_to_1mp, get_caption, draw_kps # <-- MODIFIED
 from PIL import Image
 
 class Generator:
@@ -19,16 +19,12 @@ class Generator:
         
         # Generate lineart map
         lineart_map_raw = self.mh.lineart_anime_detector(image)
-
-        # --- MODIFIED: Removed tile map ---
-        # --- 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)
-        # tile_map = tile_map_raw.resize((width, height), Image.LANCZOS) # <-- REMOVED
         
-        return depth_map, lineart_map # <-- MODIFIED
+        return depth_map, lineart_map # <-- MODIFIED (kps is now handled in predict)
 
     def predict(
         self, 
@@ -40,7 +36,6 @@ class Generator:
         img2img_strength=0.3,
         depth_strength=0.3,
         lineart_strength=0.3,
-        # tile_strength=0.7,      # <-- REMOVED
         seed=-1
     ):
         # 1. Pre-process Inputs
@@ -48,8 +43,8 @@ class Generator:
         processed_image = resize_image_to_1mp(input_image)
         target_width, target_height = processed_image.size
         
-        # 2. Get Face Embedding (Robust Mode)
-        face_emb = self.mh.get_face_embedding(processed_image)
+        # 2. Get Face Info (replaces get_face_embedding)
+        face_info = self.mh.get_face_info(processed_image)
         
         # 3. Generate Prompt
         if not user_prompt.strip():
@@ -65,43 +60,51 @@ class Generator:
         print(f"Prompt: {final_prompt}")
         print(f"Negative Prompt: {negative_prompt}")
 
-        # 4. Generate Control Maps (Structure)
-        print("Generating Control Maps (Depth, LineArt)...") # <-- MODIFIED
-        depth_map, lineart_map = self.prepare_control_images(processed_image, target_width, target_height) # <-- MODIFIED
+        # 4. Generate OTHER Control Maps (Structure)
+        print("Generating Control Maps (Depth, LineArt)...")
+        depth_map, lineart_map = self.prepare_control_images(processed_image, target_width, target_height)
         
-        # 5. Logic for Face vs No-Face
-        # --- MODIFIED: Removed Tile Control ---
-        # ControlNet order: [InstantID, Zoe, LineArt] # <-- MODIFIED
-        if face_emb is not None:
-            print("Face detected: Applying InstantID.")
-            controlnet_conditioning_scale = [0.45, depth_strength, lineart_strength] # <-- MODIFIED
-            control_guidance_end = [0.3, 0.6, 0.6] # <-- MODIFIED
-            self.mh.pipeline.set_ip_adapter_scale(0.45)
+        # 5. Logic for Face vs No-Face (NOW INCLUDES KPS)
+        # ControlNet order: [InstantID_KPS, Zoe_Depth, LineArt]
+        if face_info is not None:
+            print("Face detected: Applying InstantID with keypoints.")
+            # Get embedding
+            face_emb = torch.tensor(face_info.normed_embedding).unsqueeze(0)
+            # Create keypoint image
+            face_kps = draw_kps(processed_image, face_info['kps'])
+            
+            # Set strengths (using 0.8 from file's example)
+            controlnet_conditioning_scale = [0.8, depth_strength, lineart_strength] 
+            self.mh.pipeline.set_ip_adapter_scale(0.8)
         else:
             print("No face detected: Disabling InstantID.")
-            controlnet_conditioning_scale = [0.0, depth_strength, lineart_strength] # <-- MODIFIED
-            control_guidance_end = [0.3, 0.6, 0.6] # <-- MODIFIED
-            self.mh.pipeline.set_ip_adapter_scale(0.0)
-            
-            # --- START FIX for NoneType Error ---
+            # Create dummy embedding
             face_emb = torch.zeros((1, 512), dtype=Config.DTYPE, device=Config.DEVICE)
-            # --- END FIX ---
+            # Create dummy keypoint image (black)
+            face_kps = Image.new('RGB', (target_width, target_height), (0, 0, 0))
+            
+            # Set strengths
+            controlnet_conditioning_scale = [0.0, depth_strength, lineart_strength] 
+            self.mh.pipeline.set_ip_adapter_scale(0.0)
+
+        # We keep the guidance_end for pose low
+        control_guidance_end = [0.3, 0.6, 0.6] 
 
-        # --- ADDED: Seed/Generator Logic ---
+        # --- 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 ---
+        # --- END ---
 
         # 6. Run Inference
         print("Running pipeline...")
         result = self.mh.pipeline(
             prompt=final_prompt,
             negative_prompt=negative_prompt,
-            image=processed_image,
-            control_image=[processed_image, depth_map, lineart_map], # <-- MODIFIED
-            image_embeds=face_emb,
+            image=processed_image,  # Base img2img image
+            control_image=[face_kps, depth_map, lineart_map], # <-- MODIFIED
+            image_embeds=face_emb,  # Face identity embedding
             generator=generator,
             
             # --- Parameters from UI ---