Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers import AutoencoderKLWan, WanPipeline, UniPCMultistepScheduler
 from diffusers.utils import export_to_video
 from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_lora_to_diffusers # Keep this if it's the base
 import gradio as gr
@@ -307,6 +307,10 @@ pipe = WanPipeline.from_pretrained(
     vae=vae,
     torch_dtype=torch.bfloat16 # bfloat16 for pipeline
 )
+flow_shift = 8.0  # 5.0 for 720P, 3.0 for 480P
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=flow_shift
+)
 logger.info("Moving pipeline to CUDA...")
 pipe.to("cuda")
 
@@ -342,7 +346,7 @@ logger.info("Manual diff_b/diff patches applied.")
 
 # --- Gradio Interface Function ---
 @spaces.GPU
-def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, steps, fps):
+def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, steps, fps, progress=gr.Progress(track_tqdm=True)):
     logger.info("Starting video generation...")
     logger.info(f"  Prompt: {prompt}")
     logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
@@ -397,7 +401,7 @@ with gr.Blocks() as demo:
                 height_input = gr.Slider(minimum=256, maximum=768, step=64, value=480, label="Height (multiple of 8)")
                 width_input = gr.Slider(minimum=256, maximum=1024, step=64, value=832, label="Width (multiple of 8)")
             with gr.Row():
-                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames") # Reduced default for faster demo
+                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames") 
                 fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Output FPS")
             steps = gr.Slider(minimum=1.0, maximum=30.0, value=4.0, label="Steps")
             guidance_scale_input = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=5.0, label="Guidance Scale")