Update app.py

app.py

@@ -172,6 +172,9 @@ DEFAULT_STEPS = 15
 DEFAULT_TARGET_STEPS = 25
 # Chunked Generation: an arbitrary, adjustable starting point for the "Chunk stop (s)" field.
 DEFAULT_CHUNK_STOP = 10.0
+# Momentum: seconds of the previous chunk's tail carried into the next chunk's opening. 0 disables momentum
+# entirely, falling back to the plain last-frame-as-keyframe carry.
+DEFAULT_MOMENTUM = 2.0
 
 
 def snap_frames(seconds: float) -> int:
@@ -370,6 +373,8 @@ def _convert_diffusion_model_lora(raw: dict, base_shapes: dict, swap_fc1: bool)
     return out, network_alphas
 
 PIPE = None
+MOMENTUM_PIPE = None
+MOMENTUM_ERROR: str | None = None
 FILM = None
 FILM_ERROR: str | None = None
 LOAD_ERROR: str | None = None
@@ -385,9 +390,10 @@ def status() -> str:
         return f"Loading `{MODEL_REPO}` (transformer + VAEs, 77.3 GB). Watch the Space logs."
 
     film = "FILM **ready**" if FILM is not None else f"FILM **off** ({FILM_ERROR})"
+    momentum = "momentum **ready**" if MOMENTUM_PIPE is not None else f"momentum **off** ({MOMENTUM_ERROR})"
     return (
         f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention "
-        f"`{ATTENTION}` · {film} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
+        f"`{ATTENTION}` · {film} · {momentum} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
         f"conditioner `{CONDITIONER_SPACE}`"
     )
 
@@ -399,7 +405,7 @@ def load_models() -> str | None:
     touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio
     VAE decodes the soundtrack roughly 20 dB too quiet.
     """
-    global PIPE, FILM, FILM_ERROR, LOAD_ERROR, LOADED_IN, LORA_STATUS
+    global PIPE, FILM, FILM_ERROR, LOAD_ERROR, LOADED_IN, LORA_STATUS, MOMENTUM_PIPE, MOMENTUM_ERROR
 
     if PIPE is not None or LOAD_ERROR is not None:
         return LOAD_ERROR
@@ -513,6 +519,28 @@ def load_models() -> str | None:
             # ~10 GB of fp32 VAEs move on the first GPU call instead.
             pipe.transformer.to("cuda")
 
+        # Chunked Generation's momentum feature: a second, keyframe-free denoise graph over the *same* resident
+        # weights — `update_components` links it to `pipe`'s own `transformer`/`vae`/`audio_vae`/schedulers, so
+        # nothing here is loaded a second time. Soft-fail like FILM below: an experimental, newly-added block
+        # (`h3_momentum.py`, untested end to end) failing here shouldn't take the whole Space down with it.
+        try:
+            from h3_momentum import MiniMaxH3MomentumGeneratorBlocks
+
+            momentum_blocks = MiniMaxH3MomentumGeneratorBlocks()
+            momentum_pipe = momentum_blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
+            momentum_pipe.update_components(
+                transformer=pipe.transformer,
+                vae=pipe.vae,
+                audio_vae=pipe.audio_vae,
+                scheduler=pipe.scheduler,
+                audio_scheduler=pipe.audio_scheduler,
+            )
+            MOMENTUM_PIPE = momentum_pipe
+            print("[gen] momentum pipe ready", flush=True)
+        except Exception as error:
+            MOMENTUM_ERROR = f"{type(error).__name__}: {error}"
+            print(f"[gen] momentum pipe unavailable ({MOMENTUM_ERROR}); momentum disabled", flush=True)
+
         PIPE = pipe
         LOADED_IN = time.time() - started
         print(f"[gen] ready in {LOADED_IN:.0f}s", flush=True)
@@ -610,6 +638,12 @@ def get_duration(
     if maximize_gpu:
         return MAXIMIZE_GPU_DURATION
 
+    # Momentum: `given_video` is the second-to-last item of `*a`, matching its fixed position at the end of the
+    # `call` tuple in `generate()`. A flat, unvalidated allowance for the extra VAE encode — worth checking
+    # against a real measurement once this is testable, the same as every other constant in this function was.
+    given_video = a[-2] if len(a) >= 2 else None
+    momentum_allowance = 5 if given_video is not None else 0
+
     height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
     multiplier = max(1, int(multiplier))
     latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
@@ -626,8 +660,7 @@ def get_duration(
     film = (num_frames - 1) * (multiplier - 1) * _FILM_PER_FRAME * pixel_ratio
     post = _POST_BASE + film + out_frames * _MUX_PER_FRAME * pixel_ratio
 
-    return max(60, int((denoise + decode + post) * _MARGIN) + _PLACEMENT_ALLOWANCE)
-
+    return max(60, int((denoise + decode + post + momentum_allowance) * _MARGIN) + _PLACEMENT_ALLOWANCE)
 
 @spaces.GPU(duration=get_duration, size=GPU_SIZE)
 def _generate(
@@ -652,6 +685,8 @@ def _generate(
     stage_from,
     resume_video_latents,
     resume_audio_latents,
+    given_video,
+    video_condition_mode,
 ):
     """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,
@@ -704,13 +739,17 @@ def _generate(
     # exact same bits, every time — offsetting by `stage_from` (0 on a fresh/single-stage call, so this changes
     # nothing there) means each stage actually draws a different stream.
     effective_seed = int(seed) + stage_from
+    # Momentum: a genuinely different, keyframe-free denoise graph, sharing every weight with `PIPE` (see
+    # `load_models()`) — `pk.use_schedule`/the sampler context managers are already generic over whichever pipe
+    # object they're handed, since both read and write the same shared `scheduler`/`audio_scheduler`/`transformer`.
+    active_pipe = MOMENTUM_PIPE if given_video is not None else PIPE
     with pk.use_schedule(
-        PIPE, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=effective_seed,
+        active_pipe, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=effective_seed,
         total_steps=total_steps, stage_from=stage_from,
     ):
-        with use_dpmpp_2s_ancestral(PIPE, effective_seed, enabled=(sampler == "dpmpp_2s_ancestral")):
-            with use_dpmpp_sde_gpu(PIPE, effective_seed, enabled=(sampler == "dpmpp_sde_gpu")):
-                with use_seeds_2(PIPE, effective_seed, enabled=(sampler == "seeds_2")):
+        with use_dpmpp_2s_ancestral(active_pipe, effective_seed, enabled=(sampler == "dpmpp_2s_ancestral")):
+            with use_dpmpp_sde_gpu(active_pipe, effective_seed, enabled=(sampler == "dpmpp_sde_gpu")):
+                with use_seeds_2(active_pipe, effective_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
@@ -720,19 +759,36 @@ def _generate(
                         if resume_video_latents is not None
                         else {}
                     )
-                    state = PIPE(
-                        prompt_embeds=prompt_embeds.to("cuda"),
-                        text_token_tags=text_token_tags,
-                        image=first_frame,
-                        last_image=last_frame,
-                        height=height,
-                        width=width,
-                        num_frames=num_frames,
-                        num_inference_steps=requested_steps,
-                        output_type="pt",
-                        generator=torch.Generator("cpu").manual_seed(int(seed)),
-                        **resume_kwargs,
-                    )
+                    if given_video is not None:
+                        # Momentum: keyframe-free, so no `image`/`last_image` at all — the carried clip already
+                        # determines the opening frames more directly than a keyframe could.
+                        state = active_pipe(
+                            prompt_embeds=prompt_embeds.to("cuda"),
+                            text_token_tags=text_token_tags,
+                            height=height,
+                            width=width,
+                            num_frames=num_frames,
+                            num_inference_steps=requested_steps,
+                            output_type="pt",
+                            generator=torch.Generator("cpu").manual_seed(int(seed)),
+                            given_video=given_video.to("cuda"),
+                            video_condition_mode=video_condition_mode,
+                            **resume_kwargs,
+                        )
+                    else:
+                        state = active_pipe(
+                            prompt_embeds=prompt_embeds.to("cuda"),
+                            text_token_tags=text_token_tags,
+                            image=first_frame,
+                            last_image=last_frame,
+                            height=height,
+                            width=width,
+                            num_frames=num_frames,
+                            num_inference_steps=requested_steps,
+                            output_type="pt",
+                            generator=torch.Generator("cpu").manual_seed(int(seed)),
+                            **resume_kwargs,
+                        )
             
     denoised = time.time() - started
 
@@ -810,6 +866,7 @@ def generate(
     chunk_start=0.0,
     chunk_stop=DEFAULT_CHUNK_STOP,
     chunk_state=None,
+    momentum=DEFAULT_MOMENTUM,
     progress=gr.Progress(track_tqdm=True),
     *,
     advance: bool = False,
@@ -864,7 +921,16 @@ def generate(
     if chunk_advance and chunk_state is None:
         raise gr.Error("Press Generate with Chunked Generation enabled first, to start a chunked sequence.")
 
-    chunk_first_frame = _last_frame_path(chunk_state["paths"][-1]) if chunk_advance else None
+    chunk_first_frame = None
+    given_video = None
+    video_condition_mode = "blended"
+    if chunk_advance:
+        if MOMENTUM_PIPE is not None and float(momentum) > 0:
+            given_video = _trailing_frames(chunk_state["paths"][-1], float(momentum))
+        if given_video is None:
+            # Momentum off, unavailable, or the extraction came back empty — fall back to the plain
+            # last-frame-as-keyframe carry rather than dropping continuity entirely.
+            chunk_first_frame = _last_frame_path(chunk_state["paths"][-1])
     effective_first_frame = chunk_first_frame if chunk_advance else first_frame
     effective_last_frame = None if chunk_enabled else last_frame
 
@@ -945,6 +1011,8 @@ def generate(
         steps_done if advance else 0,
         stage_state["video_latents"] if advance else None,
         stage_state["audio_latents"] if advance else None,
+        given_video,
+        video_condition_mode,
     )
     # The same call `spaces` will book the worker with, so the report can show the fit against the measurement.
     booked_seconds = get_duration(*call)
@@ -1102,6 +1170,43 @@ def _last_frame_path(video_path: str) -> str | None:
     return path
 
 
+def _trailing_frames(video_path: str, seconds: float):
+    """The last `seconds` of `video_path`'s pixel frames, at MiniMax-H3's own native 24 fps, as `(num_frames, 3,
+    H, W)` float in `[0, 1]` — the format `MiniMaxH3MomentumConditionStep` expects for `given_video`. Strided
+    back down to 24 fps if the saved chunk was FILM-interpolated to a multiple of it: feeding the video VAE
+    frames at anything other than its own encoding rate would encode the motion at the wrong speed — a
+    real correctness point, not a cosmetic one, since FILM's own multiplier is read fresh from the UI on every
+    press. Runs on CPU; no GPU time.
+    """
+    if not video_path or seconds <= 0:
+        return None
+    import cv2
+    import numpy as np
+    import torch
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None
+    actual_fps = cap.get(cv2.CAP_PROP_FPS) or FPS
+    stride = max(1, round(actual_fps / FPS))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    start_frame = max(0, total_frames - round(seconds * actual_fps))
+
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+    frames = []
+    for index in range(total_frames - start_frame):
+        ok, frame = cap.read()
+        if not ok:
+            break
+        if index % stride == 0:
+            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        return None
+    array = np.stack(frames).astype(np.float32) / 255.0
+    return torch.from_numpy(array).permute(0, 3, 1, 2).contiguous()
+
+
 def _concat_chunks(paths: list[str]) -> str:
     """The chunks so far, concatenated with a stream copy (no re-encode) — cheap regardless of how many segments
     are in the list, so redoing the whole concat fresh on every press is simpler and more robust than trying to
@@ -1387,6 +1492,19 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
                 chunk_enabled = gr.Checkbox(label="Enable chunked generation", value=False)
                 chunk_start = gr.Number(label="Chunk start (s)", value=0.0, visible=False)
                 chunk_stop = gr.Number(label="Chunk stop (s)", value=DEFAULT_CHUNK_STOP, visible=False)
+                momentum = gr.Slider(
+                    label="Momentum (s)",
+                    minimum=0.0,
+                    maximum=5.0,
+                    step=0.5,
+                    value=DEFAULT_MOMENTUM,
+                    visible=False,
+                    info=(
+                        "Seconds of the previous chunk's tail imposed on the next chunk's opening, for real "
+                        "motion continuity — 0 falls back to a plain last-frame keyframe. Untested past a "
+                        "couple of seconds."
+                    ),
+                )
                 continue_btn = gr.Button("Continue", variant="secondary", visible=False)
 
     stage_state = gr.State(None)
@@ -1414,9 +1532,9 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
         api_name=False,
     )
     chunk_enabled.change(
-        lambda enabled: tuple(gr.update(visible=enabled) for _ in range(3)),
+        lambda enabled: tuple(gr.update(visible=enabled) for _ in range(4)),
         chunk_enabled,
-        [chunk_start, chunk_stop, continue_btn],
+        [chunk_start, chunk_stop, momentum, continue_btn],
         api_name=False,
     )
 
@@ -1457,6 +1575,7 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
         chunk_start,
         chunk_stop,
         chunk_state,
+        momentum,
     ]
 
     # `functools.partial` binds `advance`/`chunk_advance` by keyword regardless of their position in `generate`'s