Update app.py

app.py

@@ -892,13 +892,16 @@ def generate(
         for name, strength in lora_strengths.items()
         if strength > 0
     )
+    steps_done_after = steps_done + this_stage_steps
     info = [
-        f"{int(steps)} steps of `{schedule_key}`",
+        f"{this_stage_steps} steps of `{schedule_key}`",
         f"sampler `{sampler}`",
         f"shift {float(video_shift):.1f}/{float(audio_shift):.1f}",
         *post,
         f"seed {int(seed)}",
     ]
+    if stage_enabled:
+        info.append(f"staged {steps_done_after}/{int(target_steps)} steps")
     if lora_text:
         info.append(lora_text)
     report = (
@@ -906,13 +909,26 @@ def generate(
         f"{' · '.join(info)}\n\n"
         f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
         f"{', upsampled' if refined else ''}) · denoise + decode {denoise_seconds:.0f}s "
-        f"({denoise_seconds / max(1, int(steps)):.1f} s/step) · post {post_seconds:.0f}s · "
+        f"({denoise_seconds / max(1, this_stage_steps):.1f} s/step) · post {post_seconds:.0f}s · "
         f"GPU {gpu_seconds:.0f}s of {booked_seconds}s booked"
     )
     if refined:
         report += f"\n\n**Upsampled prompt**\n\n{refined}"
     print(f"[gen] {report}", flush=True)
-    return path, report
+
+    new_stage_state = (
+        {
+            "video_latents": stage_video_latents,
+            "audio_latents": stage_audio_latents,
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "steps_done": steps_done_after,
+        }
+        if stage_enabled
+        else None
+    )
+    return path, report, new_stage_state
 
 
 def _fit_keyframe(image_path, current_canvas):