Update app.py

app.py

@@ -1,69 +1,77 @@
 #!/usr/bin/env python
 """
-Gradio demo for Wan2.1 FLF2V – First & Last Frame → Video
+Gradio demo for Wan2.1-FLF2V – First & Last Frame → Video
 """
 
 import os
-# Persist HF cache on /mnt/data so it survives across launches
+
+# Persist HF cache across runs
 os.environ["HF_HOME"] = "/mnt/data/huggingface"
 
-import numpy as np
 import torch
+import numpy as np
 import gradio as gr
-from diffusers import WanImageToVideoPipeline, AutoencoderKLWan
-from diffusers.utils import export_to_video
-from transformers import CLIPVisionModel
 from PIL import Image
 import torchvision.transforms.functional as TF
+from transformers import CLIPVisionModel, CLIPImageProcessor
+from diffusers import WanImageToVideoPipeline, AutoencoderKLWan
+from diffusers.utils import export_to_video
 
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 # CONFIG
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 MODEL_ID       = "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers"
-DTYPE          = torch.float16
-MAX_AREA       = 1280 * 720
-DEFAULT_FRAMES = 81
+DTYPE          = torch.float16                # use bfloat16 if your GPU supports AMP
+MAX_AREA       = 1280 * 720                   # cap at 720p
+DEFAULT_FRAMES = 81                            # ~5s at 16fps
 
-# ----------------------------------------------------------------------
-# PIPELINE LOADING
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
+# LOAD & CACHE PIPELINE (once)
+# -----------------------------------------------------------------------------
 def load_pipeline():
-    # 1) load CLIP image encoder in full precision
-    image_encoder = CLIPVisionModel.from_pretrained(
+    # 1) CLIP vision encoder (fp32)
+    clip_encoder = CLIPVisionModel.from_pretrained(
         MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32
     )
-    # 2) load VAE in reduced precision
+    # 2) VAE in reduced precision
     vae = AutoencoderKLWan.from_pretrained(
         MODEL_ID, subfolder="vae", torch_dtype=DTYPE
     )
-    # 3) load the WanImageToVideo pipeline, balanced across GPU/CPU
+    # 3) Standard CLIPImageProcessor (needs its own config files)
+    clip_processor = CLIPImageProcessor.from_pretrained(
+        "openai/clip-vit-base-patch32",   # uses a known CLIP repo
+        use_fast=True
+    )
+    # 4) Build the Wan‐to‐video pipeline, balanced across GPU/CPU
     pipe = WanImageToVideoPipeline.from_pretrained(
         MODEL_ID,
+        image_encoder=clip_encoder,
         vae=vae,
-        image_encoder=image_encoder,
+        image_processor=clip_processor,
         torch_dtype=DTYPE,
-        device_map="balanced",            # auto-offload large modules to CPU
+        device_map="balanced",            # auto‐offload large submodules to CPU
     )
-    # 4) reduce VAE peaks & enable CPU offload for everything else
-    pipe.enable_vae_slicing()
+    # 5) Slice VAE & offload rest
+    try:
+        pipe.vae.enable_slicing()
+    except AttributeError:
+        pass
     pipe.enable_model_cpu_offload()
     return pipe
 
-# create once, at import time
 PIPE = load_pipeline()
 
-# ----------------------------------------------------------------------
-# IMAGE PREPROCESSING UTILS
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
+# IMAGE RESIZE HELPERS
+# -----------------------------------------------------------------------------
 def aspect_resize(img: Image.Image, max_area=MAX_AREA):
     ar  = img.height / img.width
-    # ensure multiple of patch size
     mod = PIPE.vae_scale_factor_spatial * PIPE.transformer.config.patch_size[1]
-    h   = round(np.sqrt(max_area * ar)) // mod * mod
-    w   = round(np.sqrt(max_area / ar)) // mod * mod
+    h   = int(np.sqrt(max_area * ar)) // mod * mod
+    w   = int(np.sqrt(max_area / ar)) // mod * mod
     return img.resize((w, h), Image.LANCZOS), h, w
 
-def center_crop_resize(img: Image.Image, h, w):
+def center_crop_resize(img: Image.Image, h: int, w: int):
     ratio = max(w / img.width, h / img.height)
     img = img.resize(
         (round(img.width * ratio), round(img.height * ratio)),
@@ -71,9 +79,9 @@ def center_crop_resize(img: Image.Image, h, w):
     )
     return TF.center_crop(img, [h, w])
 
-# ----------------------------------------------------------------------
-# GENERATION FUNCTION
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
+# GENERATION FUNCTION (streams every step)
+# -----------------------------------------------------------------------------
 def generate(
     first_frame: Image.Image,
     last_frame:  Image.Image,
@@ -86,26 +94,24 @@ def generate(
     fps:         int,
     progress=    gr.Progress()
 ):
-    # seed
+    # 1) Seed
     if seed == -1:
         seed = torch.seed()
     gen = torch.Generator(device=PIPE.device).manual_seed(seed)
 
-    # initial progress
+    # 2) Preprocess
     progress(0, steps, desc="Preprocessing images")
-
-    # resize / crop
-    first_frame, h, w = aspect_resize(first_frame)
-    if last_frame.size != first_frame.size:
+    f0, h, w = aspect_resize(first_frame)
+    if last_frame.size != f0.size:
         last_frame = center_crop_resize(last_frame, h, w)
 
-    # callback to update progress bar on each denoising step
-    def progress_callback(step, timestep, latents):
+    # 3) Streaming callback
+    def cb(step, timestep, latents):
         progress(step, steps, desc=f"Inference step {step}/{steps}")
 
-    # run the pipeline (streams progress via callback)
-    result = PIPE(
-        image=first_frame,
+    # 4) Inference
+    output = PIPE(
+        image=f0,
         last_image=last_frame,
         prompt=prompt,
         negative_prompt=negative or None,
@@ -115,46 +121,43 @@ def generate(
         num_inference_steps=steps,
         guidance_scale=guidance,
         generator=gen,
-        callback=progress_callback,
+        callback=cb
     )
 
-    # assemble and export to video
-    frames     = result.frames[0]  # list of PIL images
-    video_path = export_to_video(frames, fps=fps)
-
-    # return video and seed used (Gradio will auto-download the .mp4)
+    # 5) Export to MP4
+    video_path = export_to_video(output.frames[0], fps=fps)
     return video_path, seed
 
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 # GRADIO UI
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## Wan 2.1 FLF2V – First & Last Frame → Video")
+    gr.Markdown("## Wan2.1 FLF2V – First & Last Frame → Video")
 
     with gr.Row():
         first_img = gr.Image(label="First frame", type="pil")
         last_img  = gr.Image(label="Last frame",  type="pil")
 
-    prompt   = gr.Textbox(label="Prompt", placeholder="A blue bird takes off…")
-    negative = gr.Textbox(label="Negative prompt (optional)", placeholder="ugly, blurry")
+    prompt_box   = gr.Textbox(label="Prompt", placeholder="A blue bird takes off…")
+    negative_box = gr.Textbox(label="Negative prompt (optional)", placeholder="ugly, blurry")
 
     with gr.Accordion("Advanced parameters", open=False):
-        steps      = gr.Slider(10, 50, value=30,   step=1,   label="Sampling steps")
-        guidance   = gr.Slider(0.0, 10.0, value=5.5, step=0.1, label="Guidance scale")
-        num_frames = gr.Slider(16, 129, value=DEFAULT_FRAMES, step=1, label="Number of frames")
-        fps        = gr.Slider(4, 30, value=16,    step=1,   label="FPS (export)")
-        seed       = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
+        steps      = gr.Slider(10, 50, value=30, step=1,   label="Steps")
+        guidance   = gr.Slider(0.0, 10.0, value=5.5, step=0.1, label="Guidance")
+        num_frames = gr.Slider(16, 129, value=DEFAULT_FRAMES, step=1, label="Frames")
+        fps        = gr.Slider(4, 30, value=16, step=1,       label="FPS")
+        seed_input = gr.Number(value=-1, precision=0,        label="Seed (-1 = random)")
 
-    run_btn   = gr.Button("Generate")
-    video_out = gr.Video(label="Result (.mp4)")
-    used_seed = gr.Number(label="Seed used", interactive=False)
+    run_btn  = gr.Button("Generate")
+    video_out= gr.Video(label="Result (.mp4)")
+    seed_out = gr.Number(label="Seed used", interactive=False)
 
     run_btn.click(
         fn=generate,
-        inputs=[ first_img, last_img, prompt, negative,
-                 steps, guidance, num_frames, seed, fps ],
-        outputs=[ video_out, used_seed ]
+        inputs=[ first_img, last_img, prompt_box, negative_box,
+                 steps, guidance, num_frames, seed_input, fps ],
+        outputs=[ video_out, seed_out ]
     )
 
-# no special queue args needed
-demo.launch()
+    # no extra queue args needed; Gradio will serialize calls automatically
+    demo.launch()