Update app.py

app.py

@@ -53,9 +53,52 @@ pipeline = wan.WanTI2V(
 )
 print("Pipeline initialized and ready.")
 
+# --- Helper Functions ---
+def select_best_size_for_image(image, available_sizes):
+    """Select the size option with aspect ratio closest to the input image."""
+    if image is None:
+        return available_sizes[0]  # Return first option if no image
+    
+    img_width, img_height = image.size
+    img_aspect_ratio = img_height / img_width
+    
+    best_size = available_sizes[0]
+    best_diff = float('inf')
+    
+    for size_str in available_sizes:
+        # Parse size string like "704*1280" 
+        height, width = map(int, size_str.split('*'))
+        size_aspect_ratio = height / width
+        diff = abs(img_aspect_ratio - size_aspect_ratio)
+        
+        if diff < best_diff:
+            best_diff = diff
+            best_size = size_str
+    
+    return best_size
+
+def handle_image_upload(image):
+    """Handle image upload and return the best matching size."""
+    if image is None:
+        return gr.update()
+    
+    pil_image = Image.fromarray(image).convert("RGB")
+    available_sizes = list(SUPPORTED_SIZES[TASK_NAME])
+    best_size = select_best_size_for_image(pil_image, available_sizes)
+    
+    return gr.update(value=best_size)
+
+def get_duration(image, prompt, size, duration_seconds, sampling_steps, guide_scale, shift, seed):
+    """Calculate dynamic GPU duration based on parameters."""
+    if sampling_steps > 35 and duration_seconds > 2:
+        return 90
+    elif sampling_steps > 35 or duration_seconds > 2:
+        return 80
+    else:
+        return 60
 
 # --- 2. Gradio Inference Function ---
-@spaces.GPU(duration=80)
+@spaces.GPU(duration=get_duration)
 def generate_video(
     image,
     prompt,
@@ -71,7 +114,12 @@ def generate_video(
     if seed == -1:
         seed = random.randint(0, sys.maxsize)
 
-    input_image = Image.fromarray(image).convert("RGB") if image is not None else None
+    input_image = None
+    if image is not None:
+        input_image = Image.fromarray(image).convert("RGB")
+        # Resize image to match selected size
+        target_height, target_width = map(int, size.split('*'))
+        input_image = input_image.resize((target_width, target_height))
     
     # Calculate number of frames based on duration
     num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
@@ -134,6 +182,18 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
 
             run_button = gr.Button("Generate Video", variant="primary")
             
+    # Add image upload handler
+    image_input.upload(
+        fn=handle_image_upload,
+        inputs=[image_input],
+        outputs=[size_input]
+    )
+    
+    image_input.clear(
+        fn=handle_image_upload,
+        inputs=[image_input],
+        outputs=[size_input]
+    )
 
     example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
     gr.Examples(