Combine image+video into single GPU session to fix ZeroGPU token expiry

app.py

@@ -39,34 +39,34 @@ INPUT_DIR.mkdir(exist_ok=True)
 # GPU-accelerated steps (decorated only on Spaces)
 # ---------------------------------------------------------------------------
 
-def _generate_images(run_dir, style_name, progress_callback=None):
-    """Load SDXL pipeline, generate all images, then unload to free VRAM."""
+def _generate_images_and_videos(run_dir, style_name):
+    """GPU phase: generate all images, then all video clips.
+
+    Combined into one @spaces.GPU call so the GPU token stays valid
+    for both steps (acquiring a GPU after long CPU work expires the token).
+    """
+    # --- Images ---
     if IS_SPACES:
         from src.image_generator_hf import run as gen_images
     else:
         from src.image_generator_api import run as gen_images
-    gen_images(run_dir, style_name=style_name, progress_callback=progress_callback)
-    # Free VRAM for the video model
+    gen_images(run_dir, style_name=style_name)
     torch.cuda.empty_cache()
 
-
-def _generate_videos(run_dir, progress_callback=None):
-    """Load Wan 2.1 pipeline, generate all video clips, then unload."""
+    # --- Videos ---
     if IS_SPACES:
         from src.video_generator_hf import run as gen_videos
-        gen_videos(run_dir, progress_callback=progress_callback)
-        # Unload Wan 2.1 to free VRAM
+        gen_videos(run_dir)
         from src.video_generator_hf import unload
         unload()
     else:
         from src.video_generator_api import run as gen_videos
-        gen_videos(run_dir, progress_callback=progress_callback)
+        gen_videos(run_dir)
 
 
 # Apply @spaces.GPU decorator on Spaces
 if IS_SPACES:
-    _generate_images = spaces.GPU(duration=300)(_generate_images)
-    _generate_videos = spaces.GPU(duration=3600)(_generate_videos)  # up to 1h for ~12 clips
+    _generate_images_and_videos = spaces.GPU(duration=3600)(_generate_images_and_videos)
 
 
 # ---------------------------------------------------------------------------
@@ -323,15 +323,10 @@ def generate(audio_file: str, style_name: str, cover_art: str | None,
                          image_prompt_guidance=style.get("image_prompt_guidance", ""),
                          quality_suffix=style.get("quality_suffix", ""))
 
-    # --- Step 6: Image Generation ---
-    if step_num <= 6:
-        progress(0.50, desc="Generating images...")
-        _generate_images(run_dir, style_name)
-
-    # --- Step 7: Video Generation ---
+    # --- Steps 6 & 7: Image + Video Generation (single GPU session) ---
     if step_num <= 7:
-        progress(0.70, desc="Generating video clips...")
-        _generate_videos(run_dir)
+        progress(0.50, desc="Generating images & video clips...")
+        _generate_images_and_videos(str(run_dir), style_name)
 
     # --- Step 8: Assembly ---
     progress(0.90, desc="Assembling final video...")