Update app.py

app.py

@@ -767,6 +767,7 @@ def generate(
     stage_enabled=False,
     target_steps=DEFAULT_TARGET_STEPS,
     stage_state=None,
+    recondition=True,
     progress=gr.Progress(track_tqdm=True),
     *,
     advance: bool = False,
@@ -815,29 +816,44 @@ def generate(
     else:
         this_stage_steps = int(steps)
 
-    progress(
-        0.0,
-        desc=(
-            f"Upsampling the prompt on {CONDITIONER_SPACE} ..."
-            if upsample
-            else f"Conditioning on {CONDITIONER_SPACE} ..."
-        ),
-    )
-    conditioned = time.time()
-    prompt_embeds, text_token_tags, metadata, plan = encode_remote(
-        prompt, first_frame, last_frame, canvas, num_frames, rewrite_prompt=upsample
-    )
-    condition_seconds = time.time() - conditioned
-    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."
+    skip_recondition = advance and stage_state is not None and not recondition
+    if skip_recondition:
+        # "Re-condition" off: reuses this sequence's cached conditioning verbatim. Safe specifically because
+        # nothing sampler/schedule/shift/steps/seed/sharpen/interpolation/LoRA-related is an input to the
+        # conditioner at all — only prompt, the two keyframes, canvas, and "Upsample prompt" are. Height/width/
+        # num_frames come from that same cached conditioning, so there's nothing new to compare for the
+        # shape-consistency check below.
+        prompt_embeds = stage_state["prompt_embeds"]
+        text_token_tags = stage_state["text_token_tags"]
+        metadata = stage_state["metadata"]
+        plan = stage_state["plan"]
+        condition_seconds = 0.0
+        height, width, num_frames = stage_state["height"], stage_state["width"], stage_state["num_frames"]
+        refined = stage_state.get("refined") or ""
+    else:
+        progress(
+            0.0,
+            desc=(
+                f"Upsampling the prompt on {CONDITIONER_SPACE} ..."
+                if upsample
+                else f"Conditioning on {CONDITIONER_SPACE} ..."
+            ),
         )
+        conditioned = time.time()
+        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
+            prompt, first_frame, last_frame, canvas, num_frames, rewrite_prompt=upsample
+        )
+        condition_seconds = time.time() - conditioned
+        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
@@ -907,7 +923,7 @@ def generate(
     report = (
         f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s) -> {out_frames} frames at {fps} fps · "
         f"{' · '.join(info)}\n\n"
-        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
+        f"conditioner {condition_seconds:.0f}s{' (cached)' if skip_recondition else ''} ({plan['num_text_tokens']} tokens"
         f"{', upsampled' if refined else ''}) · denoise + decode {denoise_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"
@@ -924,6 +940,11 @@ def generate(
             "width": width,
             "num_frames": num_frames,
             "steps_done": steps_done_after,
+            "prompt_embeds": prompt_embeds,
+            "text_token_tags": text_token_tags,
+            "metadata": metadata,
+            "plan": plan,
+            "refined": refined,
         }
         if stage_enabled
         else None
@@ -1221,6 +1242,16 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
                     visible=False,
                     info="The fixed schedule's total length — 'Steps' above is how many of these one press runs.",
                 )
+                recondition = gr.Checkbox(
+                    label="Re-condition",
+                    value=True,
+                    visible=False,
+                    info=(
+                        "When turned off skips the conditioner on 'Advance' and reuses this sequence's cached prompt/keyframe "
+                        "encoding — safe as long as the prompt, keyframes, target dimension, and 'Upsample prompt' haven't "
+                        "changed since the first stage."
+                    ),
+                )
                 advance_btn = gr.Button("Advance", variant="secondary", visible=False)
 
     stage_state = gr.State(None)
@@ -1241,9 +1272,9 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
     )
 
     stage_enabled.change(
-        lambda enabled: (gr.update(visible=enabled), gr.update(visible=enabled)),
+        lambda enabled: tuple(gr.update(visible=enabled) for _ in range(3)),
         stage_enabled,
-        [target_steps, advance_btn],
+        [target_steps, recondition, advance_btn],
         api_name=False,
     )
     
@@ -1276,6 +1307,7 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
         stage_enabled,
         target_steps,
         stage_state,
+        recondition,
     ]
 
     # `functools.partial` binds `advance` by keyword regardless of its position in `generate`'s signature — the