Upload pk_workflow.py

pk_workflow.py

@@ -0,0 +1,470 @@
+"""The three things that make `Plaguekind/Minimax-H3` a *workflow* rather than just MiniMax-H3, ported onto the
+`ref2va` Space.
+
+Trimmed from the fl2va Space's `pk_workflow.py`: only the plain `scheduler.step()` samplers are here —
+`euler` (no patch at all), `euler_ancestral`, and `er_sde`. The SDE-family samplers (`dpmpp_2m_sde_gpu`,
+`dpmpp_3m_sde_gpu`) and the two-evaluation-per-step samplers (`dpmpp_2s_ancestral`, `dpmpp_sde_gpu`, `seeds_2`,
+in `h3_dpmpp_2s_ancestral.py` on the fl2va Space) aren't ported here, so neither is `torchsde` or the Brownian-tree
+noise machinery either sampler family needs.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+# ----------------------------------------------------------------------------------------------------------------
+# BasicScheduler(linear_quadratic)
+# ----------------------------------------------------------------------------------------------------------------
+# MiniMax-H3 carries two rectified-flow schedules per request, `shift = 12` for the video rows and `shift = 3` for
+# the audio rows. diffusers builds both from one `linspace(1, 0, steps)` base grid; ComfyUI instead samples the
+# *video* schedule and derives the audio one from it in closed form
+# (`comfy/ldm/minimax/model.py::time_shift_sigma`). The two agree, because the shift is a bijection of the base
+# grid — which is what lets a schedule chosen in ComfyUI's video-sigma space be transplanted here exactly.
+#
+# `linear_quadratic` is Mochi's schedule (`comfy/samplers.py::linear_quadratic_schedule`) and it does **not** go
+# through the model's shift at all: it is `sigma_max = 1.0` scaled, so the grid PlagueKind's 15 steps actually run
+# is this one verbatim, in the video stream, with the audio stream shifted off it.
+VIDEO_SHIFT = 12.0
+AUDIO_SHIFT = 3.0
+
+
+def linear_quadratic_sigmas(
+    steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None
+) -> torch.Tensor:
+    """ComfyUI's `linear_quadratic` sigma grid, in MiniMax-H3's video-sigma space.
+
+    Ported from `comfy/samplers.py::linear_quadratic_schedule` (itself from Mochi), with
+    `model_sampling.sigma_max == 1.0`, which is what a rectified-flow model has. Returns `steps + 1` strictly
+    decreasing sigmas from exactly 1.0 to exactly 0.0, so it drives `steps` forwards — ComfyUI's step count, not
+    diffusers' (where the terminal zero is one of the `num_inference_steps`).
+
+    Half the steps crawl through the first 2.5% of the trajectory and the rest sprint the remaining 97.5%: it is a
+    front-loaded schedule, which is why 15 steps of it hold up against ~28 of the native grid.
+    """
+    steps = int(steps)
+    if steps < 2:
+        return torch.tensor([1.0, 0.0], dtype=torch.float32)
+    if linear_steps is None:
+        linear_steps = steps // 2
+
+    linear = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+    threshold_noise_step_diff = linear_steps - threshold_noise * steps
+    quadratic_steps = steps - linear_steps
+    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+    const = quadratic_coef * (linear_steps**2)
+    quadratic = [quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)]
+
+    schedule = linear + quadratic + [1.0]
+    return torch.tensor([1.0 - value for value in schedule], dtype=torch.float32)
+
+
+def time_shift_sigma(sigma: torch.Tensor, from_shift: float, to_shift: float) -> torch.Tensor:
+    """Move a sigma between two exponential shifts of the same base grid.
+
+    `comfy/ldm/minimax/model.py::time_shift_sigma`: invert `sigma = s*b / (1 + (s-1)*b)` back to the base grid `b`,
+    then re-apply the other shift. Monotonic, and it fixes both 0.0 and 1.0, so a strictly decreasing schedule that
+    ends at zero stays one.
+    """
+    if from_shift == to_shift:
+        return sigma
+    base = sigma / (from_shift + sigma * (1.0 - from_shift))
+    return to_shift * base / (1.0 + (to_shift - 1.0) * base)
+
+
+# ----------------------------------------------------------------------------------------------------------------
+# BasicScheduler(sgm_uniform / simple / beta / ddim_uniform / normal)
+# ----------------------------------------------------------------------------------------------------------------
+# Five more of ComfyUI's `BasicScheduler` names, ported from `comfy/samplers.py`. Each is computed at the
+# *reference* shift (1.0 — where `time_snr_shift` is the identity, so `sigma(t) == t`) and reprojected onto each
+# scheduler's real shift by `time_shift_sigma`, exactly like `linear_quadratic_sigmas` already is and for the same
+# reason: it keeps the video and audio streams pinned to the same underlying denoising progress at each step,
+# which computing each stream's schedule independently at its own shift would not.
+#
+# `FLOW_TIMESTEPS` mirrors ComfyUI's `ModelSamplingDiscreteFlow`/`ModelSamplingAV` default of 1000 discrete steps
+# (`comfy/model_sampling.py`). Unverified specifically for MiniMax-H3's own `sampling_settings` — if a ported
+# schedule's shape looks visibly different from ComfyUI's own render at the same steps/seed, this is the first
+# thing to check.
+FLOW_TIMESTEPS = 1000
+
+
+def _reference_sigma(index_1based: int) -> float:
+    """`ModelSamplingAV.sigma(timestep)` at shift == 1.0: the shift formula is the identity, so this is just the
+    plain fraction `index / FLOW_TIMESTEPS`. `index_1based` matches ComfyUI's 1-based table construction
+    (`torch.arange(1, timesteps + 1) / timesteps`)."""
+    return index_1based / FLOW_TIMESTEPS
+
+
+def sgm_uniform_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `sgm_uniform`. Uniform in *timestep* space between the max and min sigma, dropping the point
+    that would land exactly on the minimum, then appending an exact 0.0. `steps + 1` sigmas."""
+    steps = int(steps)
+    timesteps = torch.linspace(float(FLOW_TIMESTEPS), 1.0, steps + 1)[:-1]
+    sigmas = (timesteps / FLOW_TIMESTEPS).tolist() + [0.0]
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def normal_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `normal`. Same idea as `sgm_uniform` but the linspace includes both endpoints (the minimum
+    sigma is reached exactly, not dropped), with 0.0 still appended."""
+    steps = int(steps)
+    timesteps = torch.linspace(float(FLOW_TIMESTEPS), 1.0, steps)
+    sigmas = (timesteps / FLOW_TIMESTEPS).tolist() + [0.0]
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def simple_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `simple`: evenly-spaced *indices* into the 1000-entry sigma table, walked from the high-noise
+    end, then 0.0 appended."""
+    steps = int(steps)
+    stride = FLOW_TIMESTEPS / steps
+    sigmas = [_reference_sigma(FLOW_TIMESTEPS - int(x * stride)) for x in range(steps)]
+    sigmas.append(0.0)
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def ddim_uniform_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `ddim_uniform`: a fixed-stride walk through the sigma table starting one index in, reversed so
+    the highest sigma comes first, ending at 0.0."""
+    steps = int(steps)
+    stride = max(FLOW_TIMESTEPS // steps, 1)
+    sigmas = [0.0]
+    index = 1
+    while index < FLOW_TIMESTEPS:
+        sigmas.append(_reference_sigma(index))
+        index += stride
+    sigmas.reverse()
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def beta_sigmas(steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor:
+    """ComfyUI's `beta` (arxiv.org/abs/2407.12173): table indices drawn from a Beta(alpha, beta) inverse CDF
+    instead of an even stride, biasing samples toward one end of the trajectory. Needs `scipy`."""
+    import numpy
+    import scipy.stats
+
+    steps = int(steps)
+    total = FLOW_TIMESTEPS - 1
+    positions = 1.0 - numpy.linspace(0.0, 1.0, steps, endpoint=False)
+    indices = numpy.rint(scipy.stats.beta.ppf(positions, alpha, beta) * total)
+    sigmas = []
+    last = -1
+    for value in indices:
+        if value != last:
+            sigmas.append(_reference_sigma(int(value) + 1))
+        last = value
+    sigmas.append(0.0)
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+SCHEDULE_SIGMA_FUNCS = {
+    "linear_quadratic": linear_quadratic_sigmas,
+    "sgm_uniform": sgm_uniform_sigmas,
+    "simple": simple_sigmas,
+    "beta": beta_sigmas,
+    "ddim_uniform": ddim_uniform_sigmas,
+    "normal": normal_sigmas,
+}
+
+
+def _euler_ancestral_step(scheduler, generator, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
+    """Ports k-diffusion's `sample_euler_ancestral_RF` — the flow-matching branch `sample_euler_ancestral`
+    dispatches to for `CONST`-style model sampling, which is what MiniMax-H3's `[0, 1]` sigma space is — onto one
+    `MiniMaxH3Scheduler.step()` call. Single model evaluation, same shape as `step()` itself, with fresh
+    ancestral noise injected each step instead of a plain Euler blend. Mirrors `step()`'s own care around
+    recomputing `sigma_from_timestep` from `timestep` rather than reading `self.sigmas` at the current index, for
+    the same numerical-consistency reason documented there.
+    """
+    if scheduler._step_index is None:
+        scheduler._step_index = scheduler.index_for_timestep(timestep) if scheduler._begin_index is None else scheduler._begin_index
+
+    if not isinstance(timestep, torch.Tensor):
+        timestep = torch.tensor(timestep, dtype=sample.dtype)
+    sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype)
+    while sigma_from_timestep.ndim < sample.ndim:
+        sigma_from_timestep = sigma_from_timestep.unsqueeze(-1)
+    denoised = sample + sigma_from_timestep * model_output
+
+    compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
+    sigma = scheduler.sigmas[scheduler._step_index].to(device=sample.device, dtype=compute_dtype)
+    sigma_next = scheduler.sigmas[scheduler._step_index + 1].to(device=sample.device, dtype=compute_dtype)
+    x = sample.to(dtype=compute_dtype)
+    denoised = denoised.to(dtype=compute_dtype)
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        downstep_ratio = 1 + (sigma_next / sigma - 1) * eta
+        sigma_down = sigma_next * downstep_ratio
+        alpha_next = 1 - sigma_next
+        alpha_down = 1 - sigma_down
+        renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp_min(0).sqrt()
+        ratio = sigma_down / sigma
+        prev_sample = ratio * x + (1 - ratio) * denoised
+        if eta > 0:
+            noise = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+            prev_sample = (alpha_next / alpha_down) * prev_sample + noise * s_noise * renoise_coeff
+
+    prev_sample = prev_sample.to(dtype=sample.dtype)
+    scheduler._step_index += 1
+    return prev_sample
+
+
+def _er_sde_step(scheduler, generator, model_output, timestep, sample, s_noise: float = 1.0, max_stage: int = 3):
+    """Ports k-diffusion's `sample_er_sde` (VP ER-SDE-Solver-3, arXiv:2309.06169) onto one
+    `MiniMaxH3Scheduler.step()` call. Single model evaluation per step — second/third-order accuracy comes from
+    the previous one or two steps' denoised estimates, not an extra evaluation this step — so it carries history
+    on the scheduler instance across calls, reset each request by `use_schedule` alongside `_step_index`.
+    """
+    if scheduler._step_index is None:
+        scheduler._step_index = scheduler.index_for_timestep(timestep) if scheduler._begin_index is None else scheduler._begin_index
+    i = scheduler._step_index
+
+    if not isinstance(timestep, torch.Tensor):
+        timestep = torch.tensor(timestep, dtype=sample.dtype)
+    sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype)
+    while sigma_from_timestep.ndim < sample.ndim:
+        sigma_from_timestep = sigma_from_timestep.unsqueeze(-1)
+    denoised = sample + sigma_from_timestep * model_output
+
+    compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
+    sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    sigma, sigma_next = sigmas[i], sigmas[i + 1]
+    x = sample.to(dtype=compute_dtype)
+    denoised = denoised.to(dtype=compute_dtype)
+
+    if i == 0 and float(sigma) >= 1.0:
+        # `1 - sigma` sits in a denominator below; MiniMax-H3's first sigma is exactly 1.0, so nudge it a hair
+        # under 1.0 for this sampler's math only, matching ComfyUI's `offset_first_sigma_for_snr`. Does not
+        # touch `sigma_from_timestep` above — the model was still conditioned on the real timestep.
+        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
+        shift = float(scheduler.shift)
+        sigma = shift * base / (1 + (shift - 1) * base)
+
+    def er_lambda(s):
+        return s / (1 - s)
+
+    def noise_scaler(v):
+        return v * (v**0.3).exp() + v * 10.0
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        er_lambda_s, er_lambda_t = er_lambda(sigma), er_lambda(sigma_next)
+        alpha_s, alpha_t = 1 - sigma, 1 - sigma_next
+        r_alpha = alpha_t / alpha_s
+        r = noise_scaler(er_lambda_t) / noise_scaler(er_lambda_s)
+
+        prev_sample = r_alpha * r * x + alpha_t * (1 - r) * denoised
+
+        stage_used = min(max_stage, i + 1)
+        if stage_used >= 2:
+            num_points = 200
+            dt = er_lambda_t - er_lambda_s
+            step_size = -dt / num_points
+            positions = er_lambda_t + torch.arange(num_points, device=x.device, dtype=compute_dtype) * step_size
+            scaled = noise_scaler(positions)
+
+            s_term = torch.sum(1 / scaled) * step_size
+            er_lambda_prev = er_lambda(sigmas[i - 1])
+            denoised_d = (denoised - scheduler._er_sde_old_denoised) / (er_lambda_s - er_lambda_prev)
+            prev_sample = prev_sample + alpha_t * (dt + s_term * noise_scaler(er_lambda_t)) * denoised_d
+
+            if stage_used >= 3:
+                s_u_term = torch.sum((positions - er_lambda_s) / scaled) * step_size
+                er_lambda_prev2 = er_lambda(sigmas[i - 2])
+                denoised_u = (denoised_d - scheduler._er_sde_old_denoised_d) / ((er_lambda_s - er_lambda_prev2) / 2)
+                prev_sample = prev_sample + alpha_t * ((dt**2) / 2 + s_u_term * noise_scaler(er_lambda_t)) * denoised_u
+            scheduler._er_sde_old_denoised_d = denoised_d
+
+        if s_noise > 0:
+            noise = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+            spread = (er_lambda_t**2 - er_lambda_s**2 * r**2).clamp_min(0).sqrt()
+            prev_sample = prev_sample + alpha_t * noise * s_noise * spread
+
+    scheduler._er_sde_old_denoised = denoised
+    prev_sample = prev_sample.to(dtype=sample.dtype)
+    scheduler._step_index += 1
+    return prev_sample
+
+
+class use_schedule:
+    """Set each scheduler's shift for one request, and — for anything but `native` — force its sigma grid onto
+    one of `SCHEDULE_SIGMA_FUNCS`'s named schedules.
+
+    `MiniMaxH3Scheduler.shift` is a read-only property, so a different shift means swapping in a freshly built
+    scheduler via `from_config(..., shift=...)` rather than mutating one in place — the standard diffusers idiom
+    for changing a `ConfigMixin` parameter after construction, and correct regardless of exactly how `shift` is
+    stored internally. Applied unconditionally, including under `native`, so the shift sliders affect the
+    pipeline's own default schedule too — and always restored on exit, since `pipe.scheduler`/`pipe.audio_scheduler`
+    are shared, request-spanning objects that must not carry one request's shift into the next.
+    """
+
+    def __init__(self, pipe, steps: int, schedule_name: str, video_shift: float, audio_shift: float, sampler_name: str = "euler", seed: int = 0, threshold_noise: float = 0.025):
+        self.pipe = pipe
+        self.attr_names = ["scheduler", "audio_scheduler"]
+        self.shifts = [float(video_shift), float(audio_shift)]
+        self.schedule_name = schedule_name
+        self.sampler_name = sampler_name
+        self.seed = int(seed)
+        self.steps = int(steps)
+        self.threshold_noise = float(threshold_noise)
+        self._originals: dict = {}
+
+    def __enter__(self):
+        for attr_name, shift in zip(self.attr_names, self.shifts):
+            original = getattr(self.pipe, attr_name)
+            self._originals[attr_name] = original
+            if float(original.shift) != shift:
+                setattr(self.pipe, attr_name, type(original).from_config(original.config, shift=shift))
+
+        if self.schedule_name != "native":
+            sigma_func = SCHEDULE_SIGMA_FUNCS[self.schedule_name]
+            base = sigma_func(self.steps, self.threshold_noise) if sigma_func is linear_quadratic_sigmas else sigma_func(self.steps)
+            for attr_name in self.attr_names:
+                scheduler = getattr(self.pipe, attr_name)
+                sigmas = time_shift_sigma(base, 1.0, float(scheduler.shift))
+                unbound = type(scheduler).set_timesteps
+
+                def forced(num_inference_steps=None, device=None, sigmas=None, _s=scheduler, _grid=sigmas, _f=unbound):
+                    return _f(_s, None, device, _grid)
+
+                scheduler.set_timesteps = forced
+
+        if self.sampler_name == "euler_ancestral":
+            # Separate `torch.Generator` per scheduler (offset seeds) so video and audio ancestral noise don't
+            # correlate — each generator advances across every step call to *that* scheduler over the request.
+            for offset, attr_name in enumerate(self.attr_names):
+                scheduler = getattr(self.pipe, attr_name)
+                generator = torch.Generator(device="cpu").manual_seed(self.seed + offset)
+
+                def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _g=generator, **_kwargs):
+                    return (_euler_ancestral_step(_s, _g, model_output, timestep, sample),)
+
+                scheduler.step = stepped
+        elif self.sampler_name == "er_sde":
+            for offset, attr_name in enumerate(self.attr_names):
+                scheduler = getattr(self.pipe, attr_name)
+                scheduler._er_sde_old_denoised = None
+                scheduler._er_sde_old_denoised_d = None
+                generator = torch.Generator(device="cpu").manual_seed(self.seed + offset)
+
+                def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _g=generator, **_kwargs):
+                    return (_er_sde_step(_s, _g, model_output, timestep, sample),)
+
+                scheduler.step = stepped
+        return self
+
+    def __exit__(self, *_):
+        for attr_name, original in self._originals.items():
+            current = getattr(self.pipe, attr_name)
+            current.__dict__.pop("set_timesteps", None)
+            current.__dict__.pop("step", None)
+            setattr(self.pipe, attr_name, original)
+        return False
+
+
+# ----------------------------------------------------------------------------------------------------------------
+# ImageSharpenKJ(rcas, 0.3)
+# ----------------------------------------------------------------------------------------------------------------
+def rcas(video: torch.Tensor, strength: float, chunk: int = 16) -> torch.Tensor:
+    """AMD FidelityFX **RCAS** — Robust Contrast Adaptive Sharpening — on `(frames, 3, H, W)` in `[0, 1]`.
+
+    The FidelityFX kernel, which is what `ImageSharpenKJ`'s `rcas` mode is: a 5-tap cross, a sharpening lobe whose
+    strength is limited per pixel so the ring it would create cannot leave `[0, 1]`, and a renormalised blend.
+
+        lobe = clamp(attenuation * min over channels of max(-min / 4*max, -(1 - max) / 4*(1 - min)), -0.1875, 0)
+        out  = (center + lobe * (n + s + e + w)) / (1 + 4 * lobe)
+
+    `lobe` is negative, so the neighbours are subtracted: a high-pass with a headroom-aware gain, which is why it
+    sharpens MiniMax-H3's slightly soft VAE output without haloing it. PlagueKind's 0.3 is the strength; the note in
+    the workflow calls it "very natural" and that matches — the lobe clamp caps it well below a visible ring.
+
+    Batched over `chunk` frames at a time rather than ComfyUI's one, and written back in place: the clip is already
+    resident on the card, but this runs immediately after the denoise loop's allocation peak, and a whole-clip pass at
+    the full 1344x768x124 would ask the allocator for ~8 GB of intermediates at exactly the wrong moment.
+    """
+    if strength <= 0:
+        return video
+
+    frames, _, height, width = video.shape
+    strength = float(strength)
+    for start in range(0, frames, chunk):
+        center = video[start : start + chunk]
+        padded = torch.nn.functional.pad(center, (1, 1, 1, 1), mode="reflect")
+        north = padded[:, :, 0:height, 1 : width + 1]
+        south = padded[:, :, 2 : height + 2, 1 : width + 1]
+        west = padded[:, :, 1 : height + 1, 0:width]
+        east = padded[:, :, 1 : height + 1, 2 : width + 2]
+
+        low = torch.minimum(torch.minimum(torch.minimum(torch.minimum(north, south), west), east), center)
+        high = torch.maximum(torch.maximum(torch.maximum(torch.maximum(north, south), west), east), center)
+
+        hit_min = -low / (high * 4.0 + 1e-6)
+        hit_max = -(1.0 - high) / ((1.0 - low) * 4.0 + 1e-6)
+        lobe = torch.maximum(hit_min, hit_max).amin(dim=1, keepdim=True)
+        lobe = (lobe * strength).clamp_(-0.1875, 0.0)
+        del low, high, hit_min, hit_max
+
+        neighbours = north + south + east + west
+        center.copy_(((center + lobe * neighbours) / (1.0 + 4.0 * lobe)).clamp_(0.0, 1.0))
+    return video
+
+
+# ----------------------------------------------------------------------------------------------------------------
+# FrameInterpolate(film_net_fp16, multiplier=2)
+# ----------------------------------------------------------------------------------------------------------------
+FILM_REPO = "Comfy-Org/frame_interpolation"
+FILM_FILE = "frame_interpolation/film_net_fp16.safetensors"
+
+
+def load_film():
+    """FILM, off the same checkpoint the workflow names. CPU work; `None` on any failure, and the caller skips."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    from film_net import FILMNet
+
+    path = hf_hub_download(FILM_REPO, FILM_FILE)
+    model = FILMNet()
+    model.load_state_dict(load_file(path))
+    return model.eval().to(torch.float16)
+
+
+@torch.no_grad()
+def interpolate(model, video: torch.Tensor, multiplier: int = 2) -> torch.Tensor:
+    """`multiplier`x frame interpolation of `(frames, 3, H, W)` in `[0, 1]`, FILM, on the card.
+
+    Mirrors ComfyUI's `FrameInterpolate`: one pass per adjacent pair, the flow computed once per pair and reused for
+    every intermediate timestep (`forward_multi_timestep`), and the feature pyramid of frame `i + 1` carried over as
+    frame `i` of the next pair — which halves the feature extractions. Output length is
+    `(frames - 1) * multiplier + 1`, i.e. 24 fps in, `24 * multiplier` fps out.
+    """
+    frames = video.shape[0]
+    if model is None or frames < 2 or multiplier < 2:
+        return video
+
+    dtype = torch.float16
+    timesteps = [t / multiplier for t in range(1, multiplier)]
+    # float16, not the input's float32: the buffer is the largest allocation of the whole post chain (a 2x pass over
+    # 124 frames at 1344x768 is 247 of them) and it happens right after the denoise loop's peak.
+    out = torch.empty(((frames - 1) * multiplier + 1, *video.shape[1:]), dtype=dtype, device=video.device)
+    out[0] = video[0]
+    cursor = 1
+
+    cache: dict = {}
+    for index in range(frames - 1):
+        first = video[index : index + 1].to(dtype)
+        second = video[index + 1 : index + 2].to(dtype)
+        cache["img0"] = cache.pop("next") if "next" in cache else model.extract_features(first)
+        cache["img1"] = model.extract_features(second)
+        cache["next"] = cache["img1"]
+
+        middles = model.forward_multi_timestep(first, second, timesteps, cache=cache)
+        out[cursor : cursor + len(timesteps)] = middles.to(video.dtype).clamp_(0.0, 1.0)
+        cursor += len(timesteps)
+        out[cursor] = video[index + 1]
+        cursor += 1
+
+    return out