Improve & simplify model loading

app.py

@@ -7,6 +7,7 @@ from pathlib import Path
 
 import gradio as gr
 from huggingface_hub import snapshot_download
+import spaces
 
 warnings.filterwarnings("ignore")
 
@@ -21,60 +22,27 @@ MODEL_ID = "Wan-AI/Wan2.1-FLF2V-14B-720P"
 # Global variables
 prompt_expander = None
 wan_flf2v_720P = None
-model_loaded = False
 
+# Download model snapshots from Hugging Face Hub
+print(f"Downloading/loading checkpoints for {MODEL_ID}...")
+ckpt_dir = snapshot_download(MODEL_ID, local_dir_use_symlinks=False)
+print(f"Using checkpoints from {ckpt_dir}")
 
-def load_model():
-    """Load the model from Hugging Face Hub (ZeroGPU compatible)"""
-    global wan_flf2v_720P, model_loaded
+# Load the model configuration
+cfg = WAN_CONFIGS["flf2v-14B"]
 
-    if model_loaded and wan_flf2v_720P is not None:
-        return "Model already loaded"
-
-    try:
-        gc.collect()
-
-        print(
-            "Loading Wan2.1-FLF2V-14B-720P model from Hugging Face Hub...",
-            end="",
-            flush=True,
-        )
-        cfg = WAN_CONFIGS["flf2v-14B"]
-
-        # Download model from Hugging Face Hub to local cache
-        print("\nDownloading model files from Hugging Face Hub...", flush=True)
-        # Use HF_HOME environment variable if set (for Hugging Face Spaces)
-        # Otherwise use default cache location
-        cache_base = os.environ.get("HF_HOME")
-        if cache_base:
-            cache_dir_path = Path(cache_base) / "hub"
-        else:
-            cache_dir_path = Path.home() / ".cache" / "huggingface" / "hub"
-
-        checkpoint_dir = snapshot_download(
-            repo_id=MODEL_ID,
-            cache_dir=str(cache_dir_path),
-            local_files_only=False,
-        )
-        print(f"Model downloaded to: {checkpoint_dir}", flush=True)
-
-        wan_flf2v_720P = wan.WanFLF2V(
-            config=cfg,
-            checkpoint_dir=checkpoint_dir,
-            device_id=0,
-            rank=0,
-            t5_fsdp=False,
-            dit_fsdp=False,
-            use_usp=False,
-        )
-
-        model_loaded = True
-        print(" done", flush=True)
-        return "Model loaded successfully!"
-    except Exception as e:
-        error_msg = f"Error loading model: {str(e)}"
-        print(error_msg)
-        return error_msg
+# Instantiate the model in the global scope
+print("Initializing WanFLF2V pipeline...")
+wan_flf2v_720P = wan.WanFLF2V(
+    config=cfg,
+    checkpoint_dir=ckpt_dir,
+    device_id=0,
+    rank=0,
+    t5_fsdp=False,
+    dit_fsdp=False,
+    use_usp=False,
+)
+print("Pipeline initialized and ready.")
 
 
 def prompt_enhance(prompt, img_first, img_last, tar_lang):
@@ -108,6 +76,12 @@ def prompt_enhance(prompt, img_first, img_last, tar_lang):
         return prompt
 
 
+def get_duration(sd_steps, *args, **kwargs):
+    """Calculate dynamic GPU duration based on parameters."""
+    return sd_steps * 15
+
+
+@spaces.GPU(duration=get_duration)
 def flf2v_generation(
     flf2vid_prompt,
     flf2vid_image_first,
@@ -120,11 +94,12 @@ def flf2v_generation(
     n_prompt,
     sample_solver,
     frame_num,
+    progress=gr.Progress(track_tqdm=True),
 ):
     """Generate video from first and last frame images + text prompt"""
 
     if wan_flf2v_720P is None:
-        return None, "Model is still loading. Please wait a moment and try again."
+        return None, "Model failed to load. Please check the logs."
 
     if flf2vid_image_first is None or flf2vid_image_last is None:
         return None, "Please upload both first and last frame images"
@@ -374,24 +349,11 @@ if __name__ == "__main__":
         print(" done", flush=True)
     except Exception as e:
         print(f"Warning: Could not initialize prompt expander on startup: {e}")
-        print("Prompt enhancement will be disabled until model is loaded.")
+        print("Prompt enhancement will be disabled.")
         prompt_expander = None
 
-    # Load model automatically on startup
-    print("\n" + "=" * 50)
-    print("Loading Wan2.1-FLF2V-14B-720P model...")
-    print("=" * 50)
-    load_model()
-    if wan_flf2v_720P is not None:
-        print("✓ Model loaded successfully!")
-    else:
-        print(
-            "✗ Failed to load model. The app will still start, but video generation will not work."
-        )
-    print("=" * 50 + "\n")
-
     demo = create_interface()
 
     # Launch with ZeroGPU support
-    # ZeroGPU spaces automatically handle GPU allocation
+    # ZeroGPU spaces automatically handle GPU allocation via @spaces.GPU decorator
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

requirements.txt

@@ -15,4 +15,5 @@ gradio>=5.0.0
 numpy>=1.23.5,<2
 huggingface-hub
 Pillow
+spaces
 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.6.3+cu128torch2.8-cp310-cp310-linux_x86_64.whl