Update app.py

app.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 """
 Gradio demo for Wan2.1 FLF2V – First & Last Frame → Video
-Auto-loads the fast processor and avoids missing preprocessor_config.json.
+Loads once, uses balanced device placement, streams high-level progress,
+and auto-offers the .mp4 for download.
 """
 import os
 import numpy as np
@@ -21,31 +22,29 @@ DTYPE          = torch.float16
 MAX_AREA       = 1280 * 720
 DEFAULT_FRAMES = 81
 
-# Persist cache so safetensors only download once
+# keep Hugging Face cache on disk so we don't re-download
 os.environ["HF_HOME"] = "/mnt/data/huggingface"
 
 # -----------------------------------------------------------------------------
-# LOAD PIPELINE ONCE
+# PIPELINE LOADED ONCE
 # -----------------------------------------------------------------------------
 def load_pipeline():
-    # 1) Image encoder (fp32)
+    # 1) image encoder in full precision
     image_encoder = CLIPVisionModel.from_pretrained(
         MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32
     )
-    # 2) VAE (half-precision) + slicing
+    # 2) VAE in half precision (no slicing API here)
     vae = AutoencoderKLWan.from_pretrained(
         MODEL_ID, subfolder="vae", torch_dtype=DTYPE
     )
-    vae.enable_slicing()
-
-    # 3) Pipeline, balanced across GPU & CPU, fast processor by default
+    # 3) load full pipeline balanced across GPU/CPU, with the fast processor
     pipe = WanImageToVideoPipeline.from_pretrained(
         MODEL_ID,
         image_encoder=image_encoder,
         vae=vae,
         torch_dtype=DTYPE,
-        device_map="balanced",
-        use_fast=True,              # get the fast CLIPImageProcessor internally
+        device_map="balanced",    # spreads the model to fit your 24 GB
+        use_fast=True,            # get the fast CLIPImageProcessor internally
     )
     return pipe
 
@@ -72,7 +71,7 @@ def center_crop_resize(img: Image.Image, h, w):
 
 
 # -----------------------------------------------------------------------------
-# GENERATION WITH PROGRESS STREAMING
+# GENERATE WITH STREAMING PROGRESS
 # -----------------------------------------------------------------------------
 def generate(
     first_frame: Image.Image,
@@ -84,7 +83,7 @@ def generate(
     num_frames:  int,
     seed:        int,
     fps:         int,
-    progress=    gr.Progress(),   # built-in streamer
+    progress=    gr.Progress(),
 ):
     # seed
     if seed == -1:
@@ -98,10 +97,10 @@ def generate(
         progress(0.15, desc="Resizing last frame…")
         last_resized = center_crop_resize(last_frame, h, w)
     else:
-        last_resized = first_resized  # same size
+        last_resized = first_resized
 
-    # 15–25%: setup
-    progress(0.25, desc="Launching pipeline…")
+    # 15–25%: warm up
+    progress(0.25, desc="Initializing pipeline…")
     out = PIPE(
         image=first_resized,
         last_image=last_resized,
@@ -115,8 +114,8 @@ def generate(
         generator=gen,
     )
 
-    # 25–90%: we assume the pipeline prints its own bars in console
-    progress(0.90, desc="Building video…")
+    # 25–90%: inference happens inside the pipeline (console shows bars)
+    progress(0.90, desc="Exporting video…")
     video_path = export_to_video(out.frames[0], fps=fps)
 
     # done
@@ -135,7 +134,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         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="blurry, lowres")
+    negative = gr.Textbox(label="Negative prompt (optional)", placeholder="blurry, low-res")
 
     with gr.Accordion("Advanced parameters", open=False):
         steps      = gr.Slider(10, 50, value=30, step=1, label="Steps")
@@ -155,5 +154,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[ download, seed_used ],
     )
 
-    # queue() so tasks are serialized with a top-right mini-progress indicator
+    # serialize tasks with a mini progress badge
     demo.queue().launch(server_name="0.0.0.0", server_port=7860)