Update generator.py

generator.py

@@ -1,6 +1,6 @@
 import torch
 from config import Config
-from utils import resize_image_to_1mp, get_caption, draw_kps # <-- MODIFIED
+from utils import resize_image_to_1mp, get_caption, draw_kps
 from PIL import Image
 
 class Generator:
@@ -24,18 +24,18 @@ class Generator:
         depth_map = depth_map_raw.resize((width, height), Image.LANCZOS)
         lineart_map = lineart_map_raw.resize((width, height), Image.LANCZOS)
         
-        return depth_map, lineart_map # <-- MODIFIED (kps is now handled in predict)
+        return depth_map, lineart_map
 
     def predict(
         self, 
         input_image, 
         user_prompt="",
         negative_prompt="",
-        guidance_scale=1.0,
-        num_inference_steps=10,
-        img2img_strength=0.8,
-        depth_strength=0.8,
-        lineart_strength=0.8,
+        guidance_scale=1.5,
+        num_inference_steps=6,
+        img2img_strength=0.3,
+        depth_strength=0.3,
+        lineart_strength=0.3,
         seed=-1
     ):
         # 1. Pre-process Inputs
@@ -68,8 +68,15 @@ class Generator:
         # 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)
+            
+            # --- MODIFIED: Explicitly cast tensor to Config.DTYPE ---
+            face_emb = torch.tensor(
+                face_info.normed_embedding, 
+                dtype=Config.DTYPE,  # <-- This is the fix
+                device=Config.DEVICE
+            ).unsqueeze(0)
+            # --- END MODIFIED ---
+
             # Create keypoint image
             face_kps = draw_kps(processed_image, face_info['kps'])
             
@@ -88,7 +95,7 @@ class Generator:
             self.mh.pipeline.set_ip_adapter_scale(0.0)
 
         # We keep the guidance_end for pose low
-        control_guidance_end = [0.6, 0.9, 0.9] 
+        control_guidance_end = [0.3, 0.6, 0.6] 
 
         # --- Seed/Generator Logic ---
         if seed == -1 or seed is None:
@@ -103,7 +110,7 @@ class Generator:
             prompt=final_prompt,
             negative_prompt=negative_prompt,
             image=processed_image,  # Base img2img image
-            control_image=[face_kps, depth_map, lineart_map], # <-- MODIFIED
+            control_image=[face_kps, depth_map, lineart_map],
             image_embeds=face_emb,  # Face identity embedding
             generator=generator,