Update app.py

Attempt to split the GPU token across different video clips so each has their own token.

app.py

@@ -39,36 +39,49 @@ INPUT_DIR.mkdir(exist_ok=True)
 # GPU-accelerated steps (decorated only on Spaces)
 # ---------------------------------------------------------------------------
 
-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).
-    """
-    print(f"[GPU] Entered _generate_images_and_videos(run_dir={run_dir}, style={style_name})")
-    # --- Images ---
+def _gpu_generate_images(run_dir, style_name):
+    """GPU phase: generate all images."""
+    print(f"[GPU] Generating images (run_dir={run_dir}, style={style_name})")
     if IS_SPACES:
         from src.image_generator_hf import run as gen_images
     else:
         from src.image_generator_api import run as gen_images
-    print("[GPU] Starting image generation...")
     gen_images(run_dir, style_name=style_name)
     torch.cuda.empty_cache()
+    print("[GPU] Image generation complete.")
+
+
+def _gpu_generate_one_video(run_dir, segment_idx, prompt, negative_prompt, seed):
+    """GPU phase: generate a single video clip. Each call gets a fresh ZeroGPU token."""
+    import time
+    run_dir = Path(run_dir)
+    image_path = run_dir / "images" / f"segment_{segment_idx:03d}.png"
+    clip_path = run_dir / "clips" / f"clip_{segment_idx:03d}.mp4"
+    clip_path.parent.mkdir(parents=True, exist_ok=True)
+
+    if clip_path.exists():
+        print(f"  [GPU] Clip {segment_idx}: already exists, skipping")
+        return
+
+    if not image_path.exists():
+        print(f"  [GPU] Clip {segment_idx}: image not found, skipping")
+        return
 
-    # --- Videos ---
     if IS_SPACES:
-        from src.video_generator_hf import run as gen_videos
-        gen_videos(run_dir)
-        from src.video_generator_hf import unload
-        unload()
+        from src.video_generator_hf import generate_clip
     else:
-        from src.video_generator_api import run as gen_videos
-        gen_videos(run_dir)
+        from src.video_generator_api import generate_clip
+
+    print(f"  [GPU] Generating clip {segment_idx}...")
+    t0 = time.time()
+    generate_clip(image_path, prompt, clip_path, negative_prompt, seed=seed)
+    print(f"  [GPU] Clip {segment_idx} done ({time.time() - t0:.1f}s)")
 
 
-# Apply @spaces.GPU decorator on Spaces
+# Apply @spaces.GPU decorator on Spaces — each gets a fresh token
 if IS_SPACES:
-    _generate_images_and_videos = spaces.GPU(duration=3600)(_generate_images_and_videos)
+    _gpu_generate_images = spaces.GPU(duration=300)(_gpu_generate_images)
+    _gpu_generate_one_video = spaces.GPU(duration=300)(_gpu_generate_one_video)
 
 
 # ---------------------------------------------------------------------------
@@ -216,16 +229,13 @@ _COLOR_PRESETS = {
 }
 
 
-def generate(audio_file: str, style_name: str, cover_art: str | None,
-             run_mode: str, existing_run: str | None, start_step: str | None,
-             reuse_files: bool, progress=gr.Progress()):
-    """Run the SyncAI pipeline (full or resumed).
+def generate_cpu(audio_file: str, style_name: str, cover_art: str | None,
+                 run_mode: str, existing_run: str | None, start_step: str | None,
+                 reuse_files: bool, progress=gr.Progress()):
+    """CPU phase: steps 1-5 (stems, lyrics, beats, segmentation, prompts).
 
-    Returns:
-        Path to the final video.
+    Returns state dict for the GPU phases.
     """
-    font_name = DEFAULT_FONT
-    font_color = DEFAULT_FONT_COLOR
     style = get_style(style_name)
     is_resume = run_mode == "Resume Existing"
 
@@ -241,17 +251,14 @@ def generate(audio_file: str, style_name: str, cover_art: str | None,
         print(f"Resuming {existing_run} from step {step_num}")
 
         # Always clear assembly output (cheap to redo)
-        import shutil
         out_dir = run_dir / "output"
         if out_dir.exists():
             shutil.rmtree(out_dir)
-        # Also clear intermediate assembly artifacts
         for d in ["clips_split", "clips_trimmed"]:
             p = run_dir / d
             if p.exists():
                 shutil.rmtree(p)
 
-        # If not reusing files, also clear images and video clips
         if not reuse_files:
             if step_num <= 6:
                 img_dir = run_dir / "images"
@@ -269,7 +276,6 @@ def generate(audio_file: str, style_name: str, cover_art: str | None,
     import gc
 
     def _flush_memory():
-        """Aggressively free memory between heavy ML steps."""
         gc.collect()
         if hasattr(torch, "mps") and torch.backends.mps.is_available():
             torch.mps.empty_cache()
@@ -280,7 +286,6 @@ def generate(audio_file: str, style_name: str, cover_art: str | None,
     if step_num <= 1:
         progress(0.0, desc="Separating stems...")
         from src.stem_separator import separate_stems
-        # For resume: find original audio in song dir; for new run: use uploaded file
         if is_resume:
             song_dir = run_dir.parent
             audio_candidates = list(song_dir.glob("*.wav")) + list(song_dir.glob("*.mp3")) + \
@@ -326,26 +331,63 @@ def generate(audio_file: str, style_name: str, cover_art: str | None,
                          quality_suffix=style.get("quality_suffix", ""))
         print("Prompt generation complete.")
 
-    # --- Steps 6 & 7: Image + Video Generation (single GPU session) ---
+    progress(0.45, desc="CPU steps done, requesting GPU...")
+    # Return state for GPU phases (all values must be picklable strings)
+    return str(run_dir), style_name, str(step_num), cover_art or ""
+
+
+def generate_images(run_dir_str: str, style_name: str, step_num_str: str,
+                    cover_art: str, progress=gr.Progress()):
+    """GPU phase: step 6 — generate images. Gets a fresh ZeroGPU token."""
+    step_num = int(step_num_str)
     if step_num <= 7:
-        progress(0.50, desc="Generating images & video clips...")
-        print("Starting image & video generation...")
-        try:
-            _generate_images_and_videos(str(run_dir), style_name)
-        except Exception as e:
-            print(f"ERROR in image/video generation: {type(e).__name__}: {e}")
-            import traceback
-            traceback.print_exc()
-            raise gr.Error(f"Image/video generation failed: {e}")
-
-    # --- Step 8: Assembly ---
+        progress(0.50, desc="Generating images...")
+        _gpu_generate_images(run_dir_str, style_name)
+    return run_dir_str, style_name, step_num_str, cover_art
+
+
+def generate_videos(run_dir_str: str, style_name: str, step_num_str: str,
+                    cover_art: str, progress=gr.Progress()):
+    """GPU phase: step 7 — generate video clips, one per GPU session."""
+    step_num = int(step_num_str)
+    if step_num <= 7:
+        run_dir = Path(run_dir_str)
+        with open(run_dir / "segments.json") as f:
+            segments = json.load(f)
+
+        seed = 42
+        for i, seg in enumerate(segments):
+            idx = seg["segment"]
+            prompt = seg.get("video_prompt", seg.get("scene", seg.get("prompt", "")))
+            neg = seg.get("negative_prompt", "")
+            progress(0.50 + 0.35 * (i / len(segments)),
+                     desc=f"Generating video clip {i+1}/{len(segments)}...")
+            _gpu_generate_one_video(run_dir_str, idx, prompt, neg, seed + idx)
+
+        # Unload video model after all clips are done
+        if IS_SPACES:
+            try:
+                from src.video_generator_hf import unload
+                unload()
+            except Exception:
+                pass
+
+        print(f"All {len(segments)} video clips generated.")
+    return run_dir_str, cover_art
+
+
+def generate_assembly(run_dir_str: str, cover_art: str, progress=gr.Progress()):
+    """CPU phase: step 8 — assemble final video."""
+    run_dir = Path(run_dir_str)
+    cover = cover_art if cover_art else None
+
     progress(0.90, desc="Assembling final video...")
     from src.assembler import run as assemble_video
-    final_path = assemble_video(run_dir, font_name=font_name, font_color=font_color,
-                                cover_art=cover_art)
+    final_path = assemble_video(run_dir, font_name=DEFAULT_FONT, font_color=DEFAULT_FONT_COLOR,
+                                cover_art=cover)
 
     progress(1.0, desc="Done!")
-    return str(final_path), str(run_dir), gr.update(visible=True)
+    return str(final_path), run_dir_str, gr.update(visible=True)
 
 
 def reshuffle(run_dir_str: str, cover_art: str | None, progress=gr.Progress()):
@@ -665,6 +707,12 @@ with gr.Blocks(
     reshuffle_btn = gr.Button("Reshuffle", variant="secondary", visible=False)
     last_run_dir = gr.State(value="")
 
+    # Hidden state for passing data between chained pipeline phases
+    _st_run_dir = gr.State(value="")
+    _st_style = gr.State(value="")
+    _st_step = gr.State(value="1")
+    _st_cover = gr.State(value="")
+
     # --- Event handlers ---
     example_song.change(
         fn=_on_example_song,
@@ -689,9 +737,21 @@ with gr.Blocks(
     )
 
     generate_btn.click(
-        fn=generate,
+        fn=generate_cpu,
         inputs=[audio_input, style_dropdown,
                 cover_art_input, run_mode, existing_run, start_step, reuse_files],
+        outputs=[_st_run_dir, _st_style, _st_step, _st_cover],
+    ).then(
+        fn=generate_images,
+        inputs=[_st_run_dir, _st_style, _st_step, _st_cover],
+        outputs=[_st_run_dir, _st_style, _st_step, _st_cover],
+    ).then(
+        fn=generate_videos,
+        inputs=[_st_run_dir, _st_style, _st_step, _st_cover],
+        outputs=[_st_run_dir, _st_cover],
+    ).then(
+        fn=generate_assembly,
+        inputs=[_st_run_dir, _st_cover],
         outputs=[video_output, last_run_dir, reshuffle_btn],
     )
     reshuffle_btn.click(