Update app.py

app.py

@@ -1,5 +1,3 @@
-# app.py
-
 import torch
 from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
 from diffusers.utils import export_to_video
@@ -50,8 +48,6 @@ 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):
-    if not isinstance(pil_image, Image.Image):
-        return default_h, default_w
     orig_w, orig_h = pil_image.size
     if orig_w <= 0 or orig_h <= 0:
         return default_h, default_w
@@ -69,22 +65,37 @@ def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
     
     return new_h, new_w
 
-def handle_image_upload_for_dims_wan(uploaded_pil_image):
-    new_h, new_w = _calculate_new_dimensions_wan(
-        uploaded_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)
-
-
-# NOTE: The @spaces.GPU decorator is ignored outside of the Hugging Face platform
-@spaces.GPU
-def generate_video(input_image, prompt, height, width, 
+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:
+        new_h, new_w = _calculate_new_dimensions_wan(
+            uploaded_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(input_image, prompt, height, width, 
                    negative_prompt, duration_seconds,
                    guidance_scale, steps,
                    seed, randomize_seed, 
-                   # *** KEY CHANGE 1: Remove the default value for progress ***
+                   progress):
+    if steps > 4 and duration_seconds > 2:
+        return 90
+    elif steps > 4 or duration_seconds > 2:
+        return 75
+    else:
+        return 60
+
+@spaces.GPU(duration=get_duration)
+def generate_video(input_image, prompt, height, width, 
+                   negative_prompt=default_negative_prompt, duration_seconds = 2,
+                   guidance_scale = 1, steps = 4,
+                   seed = 42, randomize_seed = False, 
                    progress=gr.Progress(track_tqdm=True)):
     
     if input_image is None:
@@ -104,8 +115,7 @@ def generate_video(input_image, prompt, height, width,
             image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
             height=target_h, width=target_w, num_frames=num_frames,
             guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
-            generator=torch.Generator(device="cuda").manual_seed(current_seed),
-            callback_on_step_end=lambda p, s, l: progress(s / steps) # Use the progress object
+            generator=torch.Generator(device="cuda").manual_seed(current_seed)
         ).frames[0]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
@@ -130,7 +140,7 @@ with gr.Blocks() as demo:
                     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_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") # Set visible=True for debugging if needed
+                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():
@@ -138,16 +148,21 @@ with gr.Blocks() as demo:
 
     input_image_component.upload(
         fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component],
+        inputs=[input_image_component, height_input, width_input],
+        outputs=[height_input, width_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]
     )
     
-    # Bundle all the UI components that the generate_video function needs
     ui_inputs = [
         input_image_component, prompt_input, height_input, width_input,
         negative_prompt_input, duration_seconds_input,
         guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
     ]
+    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
-    # *** KEY CHANGE 2: Add gr.Progress() to the list of inputs for the click event ***
-    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
+    demo.launch(share=True)