Update generator.py

generator.py

@@ -7,23 +7,31 @@ class Generator:
     def __init__(self, model_handler):
         self.mh = model_handler
 
-    def prepare_control_images(self, image, width, height):
+    def solve_bezier(self, t, p0, p1, p2, p3):
         """
-        Generates conditioning maps, ensuring they are resized
-        to the exact target dimensions (width, height).
+        Calculates a point on a cubic Bezier curve for a given t (0 to 1).
+        Formula: B(t) = (1-t)^3*P0 + 3*(1-t)^2*t*P1 + 3*(1-t)*t^2*P2 + t^3*P3
+        
+        Here, 't' is the Face Coverage Ratio.
+        The output is the Multiplier for CFG or Strength.
         """
-        print(f"Generating control maps for {width}x{height}...")
+        # Clamp t between 0 and 1 just in case
+        t = max(0.0, min(1.0, t))
         
-        # Generate depth map
-        depth_map_raw = self.mh.leres_detector(image) 
+        # Bernstein polynomials
+        term0 = (1 - t)**3 * p0
+        term1 = 3 * (1 - t)**2 * t * p1
+        term2 = 3 * (1 - t) * t**2 * p2
+        term3 = t**3 * p3
         
-        # Generate lineart map
+        return term0 + term1 + term2 + term3
+
+    def prepare_control_images(self, image, width, height):
+        print(f"Generating control maps for {width}x{height}...")
+        depth_map_raw = self.mh.leres_detector(image) 
         lineart_map_raw = self.mh.lineart_anime_detector(image)
-        
-        # 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)
-        
         return depth_map, lineart_map
 
     def predict(
@@ -46,41 +54,40 @@ class Generator:
         # 2. Get Face Info
         face_info = self.mh.get_face_info(processed_image)
         
-        # --- START ADAPTIVE PARAMETER LOGIC ---
+        # --- CUBIC BEZIER ADAPTIVE LOGIC ---
+        # Defaults (if no face detected)
         adaptive_cfg = guidance_scale
         adaptive_strength = img2img_strength
         
         if face_info is not None:
-            # Calculate Face Coverage Ratio
+            # 1. Calculate Face Coverage (t)
             bbox = face_info['bbox']
             face_area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
             total_area = target_width * target_height
-            coverage_ratio = face_area / total_area
+            coverage_ratio = face_area / total_area # This is our 't'
             
-            print(f"Face Coverage: {coverage_ratio:.2f} ({int(coverage_ratio * 12)}/12)")
+            print(f"Face Coverage: {coverage_ratio:.3f} ({int(coverage_ratio * 12)}/12)")
 
-            # Apply variance logic based on your requested thresholds
-            if coverage_ratio >= (8/12): # > 0.66 (High Coverage)
-                # Lower CFG by 5-15% (avg 10%), keep strength same
-                adaptive_cfg = guidance_scale * 0.90
-                adaptive_strength = img2img_strength * 1.0
-                print("-> High Coverage: Applying slight CFG reduction (-10%)")
-                
-            elif coverage_ratio >= (4/12): # 0.33 to 0.66 (Medium Coverage)
-                # CFG lower 20-30% (avg 25%), strength lower 5-10% (avg 7.5%)
-                adaptive_cfg = guidance_scale * 0.75
-                adaptive_strength = img2img_strength * 0.925
-                print("-> Medium Coverage: Lowering CFG (-25%) and Strength (-7.5%)")
-                
-            else: # < 0.33 (Low Coverage)
-                # CFG lower 30-40% (avg 35%), strength lower 10-15% (avg 12.5%)
-                adaptive_cfg = guidance_scale * 0.65
-                adaptive_strength = img2img_strength * 0.875
-                print("-> Low Coverage: Significantly lowering CFG (-35%) and Strength (-12.5%)")
+            # 2. Define Control Points based on your requirements
+            
+            # CFG CURVE:
+            # P0 (t=0.0): Lower by 35% (Multiplier 0.65)
+            # P3 (t=1.0): No change (Multiplier 1.0)
+            # P1, P2: Control handles to smooth the transition (Ease-In-Out)
+            cfg_mult = self.solve_bezier(coverage_ratio, 0.65, 0.70, 0.90, 1.0)
+            
+            # STRENGTH CURVE:
+            # P0 (t=0.0): Lower by 12.5% (Multiplier 0.875)
+            # P3 (t=1.0): No change (Multiplier 1.0)
+            str_mult = self.solve_bezier(coverage_ratio, 0.875, 0.90, 0.98, 1.0)
 
-        print(f"Adaptive CFG: {guidance_scale} -> {adaptive_cfg:.2f}")
-        print(f"Adaptive Strength: {img2img_strength} -> {adaptive_strength:.2f}")
-        # --- END ADAPTIVE PARAMETER LOGIC ---
+            # 3. Apply Multipliers
+            adaptive_cfg = guidance_scale * cfg_mult
+            adaptive_strength = img2img_strength * str_mult
+            
+            print(f"-> CFG Multiplier: {cfg_mult:.3f} | New CFG: {adaptive_cfg:.2f}")
+            print(f"-> Str Multiplier: {str_mult:.3f} | New Strength: {adaptive_strength:.2f}")
+        # --- END ADAPTIVE LOGIC ---
 
         # 3. Generate Prompt
         if not user_prompt.strip():
@@ -95,15 +102,15 @@ class Generator:
             
         print(f"Prompt: {final_prompt}")
 
-        # 4. Generate OTHER Control Maps
-        print("Generating Control Maps (Depth, LineArt)...")
+        # 4. Generate Control Maps
+        print("Generating Control Maps...")
         depth_map, lineart_map = self.prepare_control_images(processed_image, target_width, target_height)
         
-        # 5. Logic for Face vs No-Face
+        # 5. Face vs No-Face Setup
         if face_info is not None:
             print("Face detected: Applying InstantID with keypoints.")
             
-            # Corrected Raw Embedding Usage
+            # Use Raw Embedding (Fix)
             face_emb = torch.tensor(
                 face_info['embedding'], 
                 dtype=Config.DTYPE,
@@ -124,7 +131,6 @@ class Generator:
 
         control_guidance_end = [0.3, 0.6, 0.6] 
 
-        # --- Seed/Generator Logic ---
         if seed == -1 or seed is None:
             seed = torch.Generator().seed()
         generator = torch.Generator(device=Config.DEVICE).manual_seed(int(seed))
@@ -140,11 +146,11 @@ class Generator:
             image_embeds=face_emb,
             generator=generator,
             
-            # --- UPDATED: Use Adaptive Parameters ---
-            strength=adaptive_strength,       # <-- Uses calculated strength
-            guidance_scale=adaptive_cfg,      # <-- Uses calculated CFG
+            # --- Using Adaptive Values ---
+            strength=adaptive_strength,
+            guidance_scale=adaptive_cfg,
             num_inference_steps=num_inference_steps, 
-            # --------------------------------------
+            # ---------------------------
             
             controlnet_conditioning_scale=controlnet_conditioning_scale,
             control_guidance_end=control_guidance_end,