Update app.py

app.py

@@ -60,14 +60,14 @@ pipe.set_adapters(["causvid_lora"], adapter_weights=[1.0])
 MOD_VALUE = 32
 MOD_VALUE_H = MOD_VALUE_W = MOD_VALUE
 
-DEFAULT_H_SLIDER_VALUE = 512 
-DEFAULT_W_SLIDER_VALUE = 896 
+DEFAULT_H_SLIDER_VALUE = 512
+DEFAULT_W_SLIDER_VALUE = 896
 
 # New fixed max_area for the calculation formula
-NEW_FORMULA_MAX_AREA = float(480 * 832) 
+NEW_FORMULA_MAX_AREA = float(480 * 832)
 
 SLIDER_MIN_H = 128
-SLIDER_MAX_H = 896 
+SLIDER_MAX_H = 896
 SLIDER_MIN_W = 128
 SLIDER_MAX_W = 896
 
@@ -82,36 +82,22 @@ def _calculate_new_dimensions_wan(pil_image: Image.Image, mod_val: int, calculat
         return default_h, default_w
 
     aspect_ratio = orig_h / orig_w
-
-    # New calculation logic as per user request:
-    # height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-    # width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
-
-    # Calculate sqrt terms
+    
     sqrt_h_term = np.sqrt(calculation_max_area * aspect_ratio)
     sqrt_w_term = np.sqrt(calculation_max_area / aspect_ratio)
 
-    # Apply the formula: round(sqrt_term) then floor_division by mod_val, then multiply by mod_val
     calc_h = round(sqrt_h_term) // mod_val * mod_val
     calc_w = round(sqrt_w_term) // mod_val * mod_val
 
-    # Ensure calculated dimensions are at least mod_val (since round(...) // mod_val * mod_val can yield 0 if round(sqrt_term) < mod_val)
     calc_h = mod_val if calc_h < mod_val else calc_h
     calc_w = mod_val if calc_w < mod_val else calc_w
 
-    # Determine effective min/max dimensions from slider limits, ensuring they are multiples of mod_val.
-    # Slider min values (min_slider_h, min_slider_w) are assumed to be multiples of mod_val.
     effective_min_h = min_slider_h
     effective_min_w = min_slider_w
 
-    # Slider max values (max_slider_h, max_slider_w) might not be multiples of mod_val.
-    # The actual maximum value a slider can output is (its_max_limit // mod_val) * mod_val.
     effective_max_h_from_slider = (max_slider_h // mod_val) * mod_val
     effective_max_w_from_slider = (max_slider_w // mod_val) * mod_val
-
-    # Clip calc_h and calc_w (which are already multiples of mod_val)
-    # to the effective slider range (which are also multiples of mod_val).
-    # The results (new_h, new_w) will therefore also be multiples of mod_val.
+    
     new_h = int(np.clip(calc_h, effective_min_h, effective_max_h_from_slider))
     new_w = int(np.clip(calc_w, effective_min_w, effective_max_w_from_slider))
 
@@ -144,32 +130,43 @@ def handle_image_upload_for_dims_wan(uploaded_pil_image: Image.Image | None, cur
 # --- Gradio Interface Function ---
 @spaces.GPU
 def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
-                   height: int, width: int, num_frames: int,
-                   guidance_scale: float, steps: int, fps_for_conditioning_and_export: int,
-                   progress=gr.Progress(track_tqdm=True)):
+                   height: int, width: int, duration_seconds: float, # Changed from num_frames
+                   guidance_scale: float, steps: int,
+                   progress=gr.Progress(track_tqdm=True)): # Removed fps_for_conditioning_and_export
     if input_image is None:
         raise gr.Error("Please upload an input image.")
 
+    # Constants for frame calculation
+    FIXED_FPS = 24
+    MIN_FRAMES_MODEL = 8  # Based on original num_frames_input slider min
+    MAX_FRAMES_MODEL = 81 # Based on original num_frames_input slider max
+
     logger.info("Starting video generation...")
     logger.info(f"  Input Image: Uploaded (Original size: {input_image.size if input_image else 'N/A'})")
     logger.info(f"  Prompt: {prompt}")
     logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
     logger.info(f"  Target Output Height: {height}, Target Output Width: {width}")
-    logger.info(f"  Num Frames: {num_frames}, FPS for conditioning & export: {fps_for_conditioning_and_export}")
-    logger.info(f"  Guidance Scale: {guidance_scale}, Steps: {steps}")
 
     target_height = int(height)
     target_width = int(width)
-    num_frames = int(num_frames)
-    fps_val = int(fps_for_conditioning_and_export)
+    # duration_seconds is already float
     guidance_scale_val = float(guidance_scale)
     steps_val = int(steps)
 
+    # Calculate number of frames based on duration and fixed FPS
+    num_frames_for_pipeline = int(round(duration_seconds * FIXED_FPS))
+    # Clamp num_frames to be within model's supported range
+    num_frames_for_pipeline = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, num_frames_for_pipeline))
+    # Ensure at least MIN_FRAMES_MODEL if rounding leads to a very small number (or zero)
+    if num_frames_for_pipeline < MIN_FRAMES_MODEL:
+        num_frames_for_pipeline = MIN_FRAMES_MODEL
+
+    logger.info(f"  Duration: {duration_seconds:.1f}s, Fixed FPS (conditioning & export): {FIXED_FPS}")
+    logger.info(f"  Calculated Num Frames: {num_frames_for_pipeline} (clamped to [{MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL}])")
+    logger.info(f"  Guidance Scale: {guidance_scale_val}, Steps: {steps_val}")
+
+
     # Ensure dimensions are compatible.
-    # With the updated _calculate_new_dimensions_wan, height and width from sliders
-    # (after image upload auto-adjustment) should already be multiples of MOD_VALUE.
-    # This block acts as a safeguard if values come from direct slider interaction
-    # before an image upload, or if something unexpected happens.
     if target_height % MOD_VALUE_H != 0:
         logger.warning(f"Height {target_height} is not a multiple of {MOD_VALUE_H}. Adjusting...")
         target_height = (target_height // MOD_VALUE_H) * MOD_VALUE_H
@@ -177,7 +174,6 @@ def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
         logger.warning(f"Width {target_width} is not a multiple of {MOD_VALUE_W}. Adjusting...")
         target_width = (target_width // MOD_VALUE_W) * MOD_VALUE_W
 
-    # Ensure minimum size (should already be handled by _calculate_new_dimensions_wan and slider mins)
     target_height = max(MOD_VALUE_H, target_height if target_height > 0 else MOD_VALUE_H)
     target_width = max(MOD_VALUE_W, target_width if target_width > 0 else MOD_VALUE_W)
 
@@ -192,17 +188,16 @@ def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
             negative_prompt=negative_prompt,
             height=target_height,
             width=target_width,
-            num_frames=num_frames,
+            num_frames=num_frames_for_pipeline, # Use calculated and clamped num_frames
             guidance_scale=guidance_scale_val,
             num_inference_steps=steps_val,
-            generator=torch.Generator(device="cuda").manual_seed(0) # Consider making seed configurable
+            generator=torch.Generator(device="cuda").manual_seed(0)
         ).frames[0]
 
-    # Using a temporary file for video export
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
 
-    export_to_video(output_frames_list, video_path, fps=fps_val)
+    export_to_video(output_frames_list, video_path, fps=FIXED_FPS) # Use fixed FPS for export
     logger.info(f"Video successfully generated and saved to {video_path}")
     return video_path
 
@@ -213,19 +208,14 @@ penguin_image_url = "https://huggingface.co/datasets/huggingface/documentation-i
 
 with gr.Blocks() as demo:
     gr.Markdown(f"""
-    # Image-to-Video with Wan 2.1 I2V (14B) + CausVid LoRA
-    Powered by `diffusers` and `{MODEL_ID}`.
-    Model is loaded into memory when the app starts. This might take a few minutes.
-    Ensure you have a GPU with sufficient VRAM (e.g., ~24GB+ for these default settings).
-    Output Height and Width will be multiples of **{MOD_VALUE}**.
-    Uploading an image will suggest dimensions based on its aspect ratio and a pre-defined target pixel area ({NEW_FORMULA_MAX_AREA:.0f} pixels),
-    clamped to slider limits.
+    # Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA
     """)
     with gr.Row():
         with gr.Column():
             input_image_component = gr.Image(type="pil", label="Input Image (will be resized to target H/W)")
             prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
-
+            duration_seconds_input = gr.Slider(minimum=0.4, maximum=3.3, step=0.1, value=1.7, label="Duration (seconds)", info="The CausVid LoRA was trained on 24fps, Wan has 81 maximum frames limit, limiting the maximum to 3.3s")
+            
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt_input = gr.Textbox(
                     label="Negative Prompt (Optional)",
@@ -235,41 +225,35 @@ with gr.Blocks() as demo:
                 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})")
-                with gr.Row():
-                    num_frames_input = gr.Slider(minimum=8, maximum=81, step=1, value=41, label="Number of Frames") # Max 81 for this model
-                    fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=24, label="FPS (for conditioning & export)")
-                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps") # WanI2V is good with few steps
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale") # Low CFG usually better for I2V
+                
+                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps") 
+                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False) 
 
             generate_button = gr.Button("Generate Video", variant="primary")
 
         with gr.Column():
             video_output = gr.Video(label="Generated Video", interactive=False)
 
-    # Connect image upload to dimension auto-adjustment
     input_image_component.upload(
         fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input], # Pass current slider values for fallback on error
+        inputs=[input_image_component, height_input, width_input],
         outputs=[height_input, width_input]
     )
-    # Also trigger on clear, though handle_image_upload_for_dims_wan handles None input
     input_image_component.clear(
         fn=handle_image_upload_for_dims_wan,
         inputs=[input_image_component, height_input, width_input],
         outputs=[height_input, width_input]
     )
 
-
     inputs_for_click_and_examples = [
         input_image_component,
         prompt_input,
         negative_prompt_input,
         height_input,
         width_input,
-        num_frames_input,
+        duration_seconds_input,
         guidance_scale_input,
-        steps_slider,
-        fps_input
+        steps_slider
     ]
 
     generate_button.click(
@@ -280,13 +264,13 @@ with gr.Blocks() as demo:
 
     gr.Examples(
         examples=[
-            [penguin_image_url, "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE, 41, 1.0, 4, 24],
-            ["https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0001.jpg", "the frog jumps around", default_negative_prompt, 384, 640, 60, 1.0, 4, 24],
+            [penguin_image_url, "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, 896, 512, 2, 1.0, 4],
+            ["https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0001.jpg", "the frog jumps around", default_negative_prompt, 448, 832, 2, 1.0, 4],
         ],
-        inputs=inputs_for_click_and_examples,
+        inputs=inputs_for_click_and_examples, 
         outputs=video_output,
         fn=generate_video,
-        cache_examples="lazy" 
+        cache_examples="lazy"
     )
 
 if __name__ == "__main__":