Upload handler.py

handler.py

@@ -45,22 +45,26 @@ ALLOWED_TIMESTEPS = [1.0, 0.9937, 0.9875, 0.9812, 0.975, 0.9094, 0.725, 0.4219]
 # Check environment variable for pipeline support
 support_image_prompt = is_truthy(os.getenv("SUPPORT_INPUT_IMAGE_PROMPT"))
 
-def generate_valid_timesteps(num_steps: int, allowed_timesteps: List[float]) -> List[float]:
+def generate_valid_timesteps(num_steps: int, allowed_timesteps: List[float], start_high: bool = True) -> List[float]:
     """Generate valid timesteps by selecting from the allowed timesteps list"""
     if num_steps >= len(allowed_timesteps):
         return allowed_timesteps
     
     if num_steps == 1:
-        # For single step, use the highest timestep (most noisy)
-        return [allowed_timesteps[0]]
+        # For single step, use the highest timestep (most noisy) if start_high is True
+        return [allowed_timesteps[0] if start_high else allowed_timesteps[-1]]
     
-    # Select evenly spaced timesteps from the allowed list
-    indices = []
-    for i in range(num_steps):
-        idx = int(i * (len(allowed_timesteps) - 1) / (num_steps - 1))
-        indices.append(idx)
-    
-    return [allowed_timesteps[i] for i in indices]
+    if start_high:
+        # Select evenly spaced timesteps from the allowed list, starting from highest
+        indices = []
+        for i in range(num_steps):
+            idx = int(i * (len(allowed_timesteps) - 1) / (num_steps - 1))
+            indices.append(idx)
+        return [allowed_timesteps[i] for i in indices]
+    else:
+        # For cases where we don't need to start high (like second pass continuation)
+        # Take the last num_steps timesteps
+        return allowed_timesteps[-num_steps:]
 
 @dataclass
 class GenerationConfig:
@@ -680,13 +684,28 @@ class EndpointHandler:
                         second_pass["timesteps"] = config.second_pass_timesteps or DEFAULT_SECOND_PASS_TIMESTEPS
                     elif config.num_inference_steps != 8:
                         # Case 2: Use num_inference_steps with valid timesteps only
-                        # Split steps between passes (70% first pass, 30% second pass)
-                        first_pass_steps = max(1, int(config.num_inference_steps * 0.7))
-                        second_pass_steps = max(1, config.num_inference_steps - first_pass_steps)
-                        
-                        # Generate valid timesteps for each pass
-                        first_pass["timesteps"] = generate_valid_timesteps(first_pass_steps, ALLOWED_TIMESTEPS)
-                        second_pass["timesteps"] = generate_valid_timesteps(second_pass_steps, ALLOWED_TIMESTEPS)
+                        if config.num_inference_steps <= 4:
+                            # For very few steps, use a simple split with key timesteps
+                            if config.num_inference_steps == 1:
+                                first_pass["timesteps"] = [1.0]
+                                second_pass["timesteps"] = [0.4219]
+                            elif config.num_inference_steps == 2:
+                                first_pass["timesteps"] = [1.0]
+                                second_pass["timesteps"] = [0.9094, 0.4219] 
+                            elif config.num_inference_steps == 3:
+                                first_pass["timesteps"] = [1.0, 0.9094]
+                                second_pass["timesteps"] = [0.9094, 0.4219]
+                            else:  # 4 steps
+                                first_pass["timesteps"] = [1.0, 0.975, 0.9094]
+                                second_pass["timesteps"] = [0.9094, 0.725, 0.4219]
+                        else:
+                            # For more steps, split them properly
+                            first_pass_steps = max(1, int(config.num_inference_steps * 0.7))
+                            second_pass_steps = max(1, config.num_inference_steps - first_pass_steps)
+                            
+                            # Generate valid timesteps for each pass
+                            first_pass["timesteps"] = generate_valid_timesteps(first_pass_steps, ALLOWED_TIMESTEPS)
+                            second_pass["timesteps"] = generate_valid_timesteps(second_pass_steps, ALLOWED_TIMESTEPS)
                     else:
                         # Case 3: Use default optimized timesteps
                         first_pass["timesteps"] = DEFAULT_FIRST_PASS_TIMESTEPS