Dimensions

app.py

@@ -37,6 +37,15 @@ FIXED_FPS = 24
 MIN_FRAMES_MODEL = 8
 MAX_FRAMES_MODEL = 121 
 
+# Dimension calculation constants (from Wan 2.1 Fast demo)
+MOD_VALUE = 32
+DEFAULT_H_SLIDER_VALUE = 512
+DEFAULT_W_SLIDER_VALUE = 896
+NEW_FORMULA_MAX_AREA = 480.0 * 832.0 
+
+SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
+SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
+
 # Instantiate the pipeline in the global scope
 print("Initializing WanTI2V pipeline...")
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -55,44 +64,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
+# --- Helper Functions (from Wan 2.1 Fast demo) ---
+def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
+                                 min_slider_h, max_slider_h,
+                                 min_slider_w, max_slider_w,
+                                 default_h, default_w):
+    orig_w, orig_h = pil_image.size
+    if orig_w <= 0 or orig_h <= 0:
+        return default_h, default_w
+
+    aspect_ratio = orig_h / orig_w
     
-    return best_size
+    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
+    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
 
-def handle_image_upload(image):
-    """Handle image upload and return the best matching size."""
-    if image is None:
-        return gr.update()
+    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
+    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
     
-    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)
+    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
+    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
     
-    return gr.update(value=best_size)
+    return new_h, new_w
+
+def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
+    if uploaded_pil_image is None:
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+    try:
+        # Convert numpy array to PIL Image if needed
+        if hasattr(uploaded_pil_image, 'shape'):  # numpy array
+            pil_image = Image.fromarray(uploaded_pil_image).convert("RGB")
+        else:  # already PIL Image
+            pil_image = uploaded_pil_image
+            
+        new_h, new_w = _calculate_new_dimensions_wan(
+            pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
+            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
+            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
+        )
+        return gr.update(value=new_h), gr.update(value=new_w)
+    except Exception as e:
+        gr.Warning("Error attempting to calculate new dimensions")
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
 
 def get_duration(image, 
                  prompt, 
-                 size, 
+                 height,
+                 width,
                  duration_seconds, 
                  sampling_steps, 
                  guide_scale, 
@@ -112,7 +129,8 @@ def get_duration(image,
 def generate_video(
     image,
     prompt,
-    size,
+    height,
+    width,
     duration_seconds,
     sampling_steps,
     guide_scale,
@@ -124,21 +142,27 @@ def generate_video(
     if seed == -1:
         seed = random.randint(0, sys.maxsize)
 
+    # Ensure dimensions are multiples of MOD_VALUE
+    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
+    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
+
     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))
+        # Resize image to match target dimensions
+        input_image = input_image.resize((target_w, target_h))
     
     # Calculate number of frames based on duration
     num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
 
+    # Create size string for the pipeline
+    size_str = f"{target_h}*{target_w}"
+
     video_tensor = pipeline.generate(
         input_prompt=prompt,
         img=input_image,  # Pass None for T2V, Image for I2V
-        size=SIZE_CONFIGS[size],
-        max_area=MAX_AREA_CONFIGS[size],
+        size=SIZE_CONFIGS.get(size_str, (target_h, target_w)),
+        max_area=MAX_AREA_CONFIGS.get(size_str, target_h * target_w),
         frame_num=num_frames,  # Use calculated frames instead of cfg.frame_num
         shift=shift,
         sample_solver='unipc',
@@ -170,7 +194,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
 
     with gr.Row():
         with gr.Column(scale=2):
-            image_input = gr.Image(type="numpy", label="Input Image (Optional)", elem_id="input_image")
+            image_input = gr.Image(type="numpy", label="Input Image (Optional, auto-resized to target H/W)", elem_id="input_image")
             prompt_input = gr.Textbox(label="Prompt", value="A beautiful waterfall in a lush jungle, cinematic.", lines=3)
             duration_input = gr.Slider(
                 minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1), 
@@ -180,40 +204,41 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
                 label="Duration (seconds)", 
                 info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
             )
-            size_input = gr.Dropdown(label="Output Resolution", choices=list(SUPPORTED_SIZES[TASK_NAME]), value="704*1280")
-        with gr.Column(scale=2):
-            video_output = gr.Video(label="Generated Video", elem_id="output_video")
             
-
             with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
+                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
                 steps_input = gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=38, step=1)
                 scale_input = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=cfg.sample_guide_scale, step=0.1)
                 shift_input = gr.Slider(label="Sample Shift", minimum=1.0, maximum=20.0, value=cfg.sample_shift, step=0.1)
                 seed_input = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
 
+        with gr.Column(scale=2):
+            video_output = gr.Video(label="Generated Video", elem_id="output_video")
             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]
+        fn=handle_image_upload_for_dims_wan,
+        inputs=[image_input, height_input, width_input],
+        outputs=[height_input, width_input]
     )
     
     image_input.clear(
-        fn=handle_image_upload,
-        inputs=[image_input],
-        outputs=[size_input]
+        fn=handle_image_upload_for_dims_wan,
+        inputs=[image_input, height_input, width_input],
+        outputs=[height_input, width_input]
     )
 
     example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
     gr.Examples(
         examples=[
-            [example_image_path, "The cat removes the glasses from its eyes.", "1280*704", 1.5],
-            [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", "1280*704", 2.0],
-            [None, "Drone footage flying over a futuristic city with flying cars.", "1280*704", 2.0],
+            [example_image_path, "The cat removes the glasses from its eyes.", 704, 1280, 1.5],
+            [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", 704, 1280, 2.0],
+            [None, "Drone footage flying over a futuristic city with flying cars.", 704, 1280, 2.0],
         ],
-        inputs=[image_input, prompt_input, size_input, duration_input],
+        inputs=[image_input, prompt_input, height_input, width_input, duration_input],
         outputs=video_output,
         fn=generate_video,
         cache_examples=False,
@@ -221,7 +246,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
 
     run_button.click(
         fn=generate_video,
-        inputs=[image_input, prompt_input, size_input, duration_input, steps_input, scale_input, shift_input, seed_input],
+        inputs=[image_input, prompt_input, height_input, width_input, duration_input, steps_input, scale_input, shift_input, seed_input],
         outputs=video_output
     )