Update app.py

app.py

@@ -11,6 +11,7 @@ transformer and the two autoencoders. `prompt_embeds` + `text_token_tags` is the
 
 from __future__ import annotations
 
+import functools
 import os
 import tempfile
 import time
@@ -154,6 +155,9 @@ INTERPOLATION = {"off · 24 fps": 1, "2x · 48 fps (PlagueKind)": 2, "4x · 96 f
 DEFAULT_INTERPOLATION = "2x · 48 fps (PlagueKind)"
 DEFAULT_SHARPEN = 0.3
 DEFAULT_STEPS = 15
+# Staged Denoising: an arbitrary, adjustable starting point for the "Target total steps" slider — the total the
+# fixed schedule is built at, walked across however many "Advance" presses it takes at "Steps" steps per press.
+DEFAULT_TARGET_STEPS = 25
 
 
 def snap_frames(seconds: float) -> int:
@@ -615,6 +619,10 @@ def _generate(
     video_shift,
     audio_shift,
     sampler,
+    total_steps,
+    stage_from,
+    resume_video_latents,
+    resume_audio_latents,
 ):
     """The only thing on GPU time: the denoise loop, the two decoders and the workflow's post chain.
     The mp4 is muxed here rather than in the caller: a `@spaces.GPU` return crosses a process boundary by pickling,
@@ -661,10 +669,22 @@ def _generate(
     requested_steps = steps if custom_schedule else steps + 1
 
     started = time.time()
-    with pk.use_schedule(PIPE, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=int(seed)):
+    with pk.use_schedule(
+        PIPE, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=int(seed),
+        total_steps=total_steps, stage_from=stage_from,
+    ):
         with use_dpmpp_2s_ancestral(PIPE, int(seed), enabled=(sampler == "dpmpp_2s_ancestral")):
             with use_dpmpp_sde_gpu(PIPE, int(seed), enabled=(sampler == "dpmpp_sde_gpu")):
                 with use_seeds_2(PIPE, int(seed), enabled=(sampler == "seeds_2")):
+                    # Staged Denoising: resuming hands the pipeline the previous stage's own latents instead of
+                    # letting `PrepareLatentsStep` draw fresh noise — both are declared-optional inputs on that
+                    # step precisely for this ("used instead of the draw"), so nothing else about the call
+                    # changes. `resume_video_latents is None` is exactly the unstaged, fresh-start case.
+                    resume_kwargs = (
+                        {"latents": resume_video_latents.to("cuda"), "audio_latents": resume_audio_latents.to("cuda")}
+                        if resume_video_latents is not None
+                        else {}
+                    )
                     state = PIPE(
                         prompt_embeds=prompt_embeds.to("cuda"),
                         text_token_tags=text_token_tags,
@@ -676,6 +696,7 @@ def _generate(
                         num_inference_steps=requested_steps,
                         output_type="pt",
                         generator=torch.Generator("cpu").manual_seed(int(seed)),
+                        **resume_kwargs,
                     )
             
     denoised = time.time() - started
@@ -683,6 +704,10 @@ def _generate(
     video = state.get("videos")[0]  # (frames, 3, H, W), float in [0, 1], on the card
     audio = state.get("audio")[0].cpu()
     sampling_rate = state.get("sampling_rate")
+    # Staged Denoising: this stage's own final latents, ahead of decode — the state a later "Advance" press
+    # resumes from. Computed unconditionally; harmless and cheap when staging isn't in use.
+    stage_video_latents = state.get("latents").cpu()
+    stage_audio_latents = state.get("audio_latents").cpu()
     del state
     # The post chain runs on the allocator the denoise loop just left fragmented (78.5 GiB at the full canvas), and
     # RCAS and FILM both want a few contiguous gigabytes.
@@ -707,7 +732,10 @@ def _generate(
     encode_video(frames, fps=fps, output_path=path, audio=audio, audio_sample_rate=sampling_rate)
 
     # `booked` to here is what `get_duration` had to predict, so it is what the report prints it against.
-    return path, denoised, post_seconds, time.time() - booked, int(frames.shape[0]), fps, multiplier
+    return (
+        path, denoised, post_seconds, time.time() - booked, int(frames.shape[0]), fps, multiplier,
+        stage_video_latents, stage_audio_latents,
+    )
 
 
 def generate(
@@ -736,10 +764,23 @@ def generate(
     video_shift=DEFAULT_VIDEO_SHIFT,
     audio_shift=DEFAULT_AUDIO_SHIFT,
     sampler=DEFAULT_SAMPLER,
+    stage_enabled=False,
+    target_steps=DEFAULT_TARGET_STEPS,
+    stage_state=None,
     progress=gr.Progress(track_tqdm=True),
+    *,
+    advance: bool = False,
 ):
     """One request through the PlagueKind graph. Every parameter but the prompt carries the default its UI
-    component carries, so an example that fills only `prompt` (and `canvas`) behaves exactly like the button."""
+    component carries, so an example that fills only `prompt` (and `canvas`) behaves exactly like the button.
+
+    `advance` isn't a UI control — it's bound per-button via `functools.partial` (`False` for "Generate", `True`
+    for "Advance") so the two share this one function rather than duplicating the conditioning/report logic.
+    Staged Denoising, debugging-only, unlocked: nothing here stops the prompt, canvas, sampler, schedule, or
+    shift from changing between an "Advance" press and the stage before it — the only samplers actually reasoned
+    through for exact-vs-different-but-equal-quality resume behavior are `euler`, `euler_ancestral`, `seeds_2`,
+    and `dpmpp_2s_ancestral`; the SDE-family samplers are untested here and not recommended.
+    """
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
     if PIPE is None:
@@ -754,6 +795,26 @@ def generate(
     multiplier = INTERPOLATION.get(interpolation, 2)
     num_frames = snap_frames(duration)
 
+    if stage_enabled and schedule_key == "native":
+        raise gr.Error(
+            "Staged Denoising needs a named sigma schedule, not `native` — the stage boundary is a slice of a "
+            "schedule this Space builds itself, and the pipeline's own default schedule isn't one this Space "
+            "controls the construction of."
+        )
+    if advance and stage_state is None:
+        raise gr.Error("Press Generate with Staged Denoising enabled first, to start a staged sequence.")
+    steps_done = int(stage_state["steps_done"]) if (advance and stage_state) else 0
+    if advance:
+        remaining = int(target_steps) - steps_done
+        if remaining <= 0:
+            raise gr.Error(
+                f"Already at or past the target step count ({steps_done}/{int(target_steps)}). Raise "
+                f"'Target total steps' to continue."
+            )
+        this_stage_steps = min(int(steps), remaining)
+    else:
+        this_stage_steps = int(steps)
+
     progress(
         0.0,
         desc=(
@@ -770,6 +831,14 @@ def generate(
     height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
     refined = plan.get("refined_prompt") or ""
 
+    if advance and (height, width, num_frames) != (
+        int(stage_state["height"]), int(stage_state["width"]), int(stage_state["num_frames"])
+    ):
+        raise gr.Error(
+            "Canvas or duration resolved differently than the staged sequence's first stage — both have to stay "
+            "fixed across a staged sequence, since they determine the saved latents' shape."
+        )
+
     def keyframe(path):
         # The conditioning latents encoded here have to be of the image the conditioner looked at, which it
         # prepares exactly this way.
@@ -779,7 +848,7 @@ def generate(
     # needs wiring in; `_generate`, `set_adapters`, and the report line below are all keyed off this dict.
     lora_strengths = {"lora1": float(lora_1_strength), "lorah": float(lora_h_strength), "lorai": float(lora_i_strength), "loraa": float(lora_a_strength), "lorab": float(lora_b_strength), "lorac": float(lora_c_strength), "lorad": float(lora_d_strength), "lorae": float(lora_e_strength), "loraf": float(lora_f_strength), "lorag": float(lora_g_strength)}
 
-    progress(0.1, desc=f"Denoising {int(steps)} steps at {width}x{height}, {num_frames} frames ...")
+    progress(0.1, desc=f"Denoising {this_stage_steps} steps at {width}x{height}, {num_frames} frames ...")
     call = (
         prompt_embeds,
         text_token_tags,
@@ -788,7 +857,7 @@ def generate(
         height,
         width,
         num_frames,
-        int(steps),
+        this_stage_steps,
         schedule_key,
         float(sharpen),
         multiplier,
@@ -798,6 +867,10 @@ def generate(
         float(video_shift),
         float(audio_shift),
         SAMPLERS.get(sampler, "euler"),
+        int(target_steps) if stage_enabled else None,
+        steps_done if advance else 0,
+        stage_state["video_latents"] if advance else None,
+        stage_state["audio_latents"] if advance else None,
     )
     # The same call `spaces` will book the worker with, so the report can show the fit against the measurement.
     booked_seconds = get_duration(*call)
@@ -1112,6 +1185,27 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
                     value=DEFAULT_LORA_G_STRENGTH,
                 )
 
+            with gr.Accordion("Staged Denoising", open=False):
+                gr.Markdown(
+                    "**Debugging feature — not for the SDE-family samplers** (`dpmpp_2m_sde_gpu`, "
+                    "`dpmpp_3m_sde_gpu`, `dpmpp_sde_gpu`). Splits one long denoise into several cheaper requests: "
+                    "run the first stage with **Generate**, then **Advance** to keep denoising the same latents "
+                    "further, as many times as needed to reach the target."
+                )
+                stage_enabled = gr.Checkbox(label="Enable staged denoising", value=False)
+                target_steps = gr.Slider(
+                    label="Target total steps",
+                    minimum=4,
+                    maximum=100,
+                    step=1,
+                    value=DEFAULT_TARGET_STEPS,
+                    visible=False,
+                    info="The fixed schedule's total length — 'Steps' above is how many of these one press runs.",
+                )
+                advance_btn = gr.Button("Advance", variant="secondary", visible=False)
+
+    stage_state = gr.State(None)
+
     first_frame.upload(_fit_keyframe, [first_frame, canvas], [first_frame, canvas])
     last_frame.upload(_fit_keyframe, [last_frame, canvas], [last_frame, canvas])
 
@@ -1126,6 +1220,13 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
     last_frame_timestamp.change(
         _extract_frame, [video, last_frame_timestamp], last_frame, show_progress="hidden"
     )
+
+    stage_enabled.change(
+        lambda enabled: (gr.update(visible=enabled), gr.update(visible=enabled)),
+        stage_enabled,
+        [target_steps, advance_btn],
+        api_name=False,
+    )
     
     controls = [
         prompt,
@@ -1153,9 +1254,20 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
         video_shift,
         audio_shift,
         sampler,
+        stage_enabled,
+        target_steps,
+        stage_state,
     ]
 
-    run.click(generate, controls, [video, report], api_name="generate")
+    # `functools.partial` binds `advance` by keyword regardless of its position in `generate`'s signature — the
+    # two buttons share every other line of conditioning/report logic and differ only in this one flag.
+    run.click(
+        functools.partial(generate, advance=False), controls, [video, report, stage_state], api_name="generate"
+    )
+    advance_btn.click(
+        functools.partial(generate, advance=True), controls, [video, report, stage_state],
+        api_name="generate_advance",
+    )
 
 
 if __name__ == "__main__":