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
+from utils import get_caption, draw_kps
 from PIL import Image
 
 class Generator:
@@ -10,21 +10,74 @@ class Generator:
     def solve_bezier(self, t, p0, p1, p2, p3):
         """
         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
         """
         t = max(0.0, min(1.0, t))
-        
         term0 = (1 - t)**3 * p0
         term1 = 3 * (1 - t)**2 * t * p1
         term2 = 3 * (1 - t) * t**2 * p2
         term3 = t**3 * p3
-        
         return term0 + term1 + term2 + term3
 
+    def smart_crop_and_resize(self, image):
+        """
+        Analyzes aspect ratio and snaps to the best SDXL resolution bucket.
+        Performs a center crop to match the target ratio, then resizes.
+        """
+        w, h = image.size
+        aspect_ratio = w / h
+        
+        # 1. Determine Target Resolution (Horizon SDXL Buckets)
+        if 0.85 <= aspect_ratio <= 1.15:
+            # Square-ish -> 1024x1024
+            target_w, target_h = 1024, 1024
+            print(f"Snap to Bucket: Square (1024x1024)")
+            
+        elif aspect_ratio < 0.85:
+            # Portrait
+            # Decide between 896x1152 (AR ~0.77) and 832x1216 (AR ~0.68)
+            if aspect_ratio < 0.72:
+                target_w, target_h = 832, 1216 # Tall Portrait
+                print(f"Snap to Bucket: Tall Portrait (832x1216)")
+            else:
+                target_w, target_h = 896, 1152 # Standard Portrait
+                print(f"Snap to Bucket: Portrait (896x1152)")
+                
+        else: # aspect_ratio > 1.15
+            # Landscape
+            # Decide between 1152x896 (AR ~1.28) and 1216x832 (AR ~1.46)
+            if aspect_ratio > 1.35:
+                target_w, target_h = 1216, 832 # Wide Landscape
+                print(f"Snap to Bucket: Wide Landscape (1216x832)")
+            else:
+                target_w, target_h = 1152, 896 # Standard Landscape
+                print(f"Snap to Bucket: Landscape (1152x896)")
+                
+        # 2. Center Crop to Target Aspect Ratio
+        target_ar = target_w / target_h
+        
+        if aspect_ratio > target_ar:
+            # Image is wider than target -> Crop width (cut sides)
+            new_w = int(h * target_ar)
+            offset = (w - new_w) // 2
+            crop_box = (offset, 0, offset + new_w, h)
+        else:
+            # Image is taller than target -> Crop height (cut top/bottom)
+            new_h = int(w / target_ar)
+            offset = (h - new_h) // 2
+            crop_box = (0, offset, w, offset + new_h)
+            
+        cropped_img = image.crop(crop_box)
+        
+        # 3. Resize to Exact Target Resolution
+        final_img = cropped_img.resize((target_w, target_h), Image.LANCZOS)
+        return final_img
+
     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)
+        
+        # Maps are resized to match the exact bucket 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
@@ -41,9 +94,9 @@ class Generator:
         lineart_strength=0.3,
         seed=-1
     ):
-        # 1. Pre-process Inputs
+        # 1. Pre-process Inputs (New Smart Crop)
         print("Processing Input...")
-        processed_image = resize_image_to_1mp(input_image)
+        processed_image = self.smart_crop_and_resize(input_image)
         target_width, target_height = processed_image.size
         
         # 2. Get Face Info
@@ -62,15 +115,11 @@ class Generator:
             
             print(f"Face Coverage: {coverage_ratio:.3f} ({int(coverage_ratio * 12)}/12)")
 
-            # 2. Define Control Points (LESS AGGRESSIVE REDUCTION)
-            
-            # CFG CURVE:
-            # Old P0 was 0.65 (35% drop). New P0 is 0.825 (17.5% drop).
-            # Curve eases from 0.825 up to 1.0 smoothly.
+            # 2. Define Control Points (Half Less Aggressive)
+            # CFG: 0.825 start (17.5% reduction)
             cfg_mult = self.solve_bezier(coverage_ratio, 0.825, 0.85, 0.95, 1.0)
             
-            # STRENGTH CURVE:
-            # Old P0 was 0.875 (12.5% drop). New P0 is 0.9375 (~6% drop).
+            # Strength: 0.9375 start (6.25% reduction)
             str_mult = self.solve_bezier(coverage_ratio, 0.9375, 0.95, 0.99, 1.0)
 
             # 3. Apply Multipliers
@@ -102,7 +151,7 @@ class Generator:
         if face_info is not None:
             print("Face detected: Applying InstantID with keypoints.")
             
-            # Use Raw Embedding (Fix)
+            # Use Raw Embedding
             face_emb = torch.tensor(
                 face_info['embedding'], 
                 dtype=Config.DTYPE,
@@ -111,8 +160,8 @@ class Generator:
 
             face_kps = draw_kps(processed_image, face_info['kps'])
             
-            controlnet_conditioning_scale = [0.777, depth_strength, lineart_strength] 
-            self.mh.pipeline.set_ip_adapter_scale(0.777)
+            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.")
             face_emb = torch.zeros((1, 512), dtype=Config.DTYPE, device=Config.DEVICE)
@@ -121,7 +170,7 @@ class Generator:
             controlnet_conditioning_scale = [0.0, depth_strength, lineart_strength] 
             self.mh.pipeline.set_ip_adapter_scale(0.0)
 
-        control_guidance_end = [0.333, 0.666, 0.666] 
+        control_guidance_end = [0.3, 0.6, 0.6] 
 
         if seed == -1 or seed is None:
             seed = torch.Generator().seed()