Pipeline API

def setup(self, **kwargs: Any) -> Any:
    """Instantiate the configured object."""
    return self._target(self, **kwargs)

def initialize_cache(
    self,
    transformer_context: dict[str, Any] | None = None,
    encoder_context: dict[str, Any] | None = None,
    decoder_context: dict[str, Any] | None = None,
) -> StreamInferencePipelineCache[
    StreamingEncoderCacheT, TransformerCacheT, StreamingDecoderCacheT
]:
    """Build a fresh per-rollout cache.

    Each ``*_context`` dict is forwarded as keyword arguments to the
    corresponding component's ``initialize_autoregressive_cache``.

    Args:
        transformer_context: Per-rollout state for the transformer
            (e.g. ``{"text_embeddings": ..., "image_embeddings": ...}``).
        encoder_context: Per-rollout state for the encoder. Ignored
            when there is no encoder.
        decoder_context: Per-rollout state for the decoder. Ignored
            when there is no decoder.

    Returns:
        A fresh cache to thread through ``generate`` / ``finalize``.
    """
    transformer_context = transformer_context or {}
    encoder_context = encoder_context or {}
    decoder_context = decoder_context or {}
    return StreamInferencePipelineCache(
        encoder_cache=(
            self.encoder.initialize_autoregressive_cache(**encoder_context)
            if self.encoder is not None
            else None
        ),
        decoder_cache=(
            self.decoder.initialize_autoregressive_cache(**decoder_context)
            if self.decoder is not None
            else None
        ),
        transformer_cache=self.diffusion_model.transformer.initialize_autoregressive_cache(
            **transformer_context
        ),
    )

@torch.no_grad()
def generate(
    self,
    autoregressive_index: int,
    cache: StreamInferencePipelineCache[
        StreamingEncoderCacheT, TransformerCacheT, StreamingDecoderCacheT
    ],
    input: Any = None,
) -> Tensor:
    """Generate one chunk for this AR step.

    Args:
        autoregressive_index: Must be ``cache.autoregressive_index + 1``,
            or ``0`` for the first call after ``initialize_cache``.
        cache: Per-rollout cache from ``initialize_cache``.
        input: Raw input fed to the encoder. Required when an encoder
            is configured, must be ``None`` otherwise. Use
            ``NullEncoderConfig`` to pass an already-encoded tensor
            straight through.

    Returns:
        Decoded tensor (e.g. RGB video) when a decoder is configured;
        otherwise the unpatchified clean latent from the diffusion model.
    """
    return self._generate(
        autoregressive_index,
        cache,
        input=input,
    )

@torch.no_grad()
def finalize(
    self,
    autoregressive_index: int,
    cache: StreamInferencePipelineCache[
        StreamingEncoderCacheT, TransformerCacheT, StreamingDecoderCacheT
    ],
) -> dict[str, float] | None:
    """Advance the diffusion AR cache for the next AR step.

    Args:
        autoregressive_index: Must match the index passed to the most
            recent ``generate`` (asserted).
        cache: Same cache used by ``generate``. Consumes
            ``cache.final_state``.

    Returns:
        ``None`` unless stages were timed, either by the injected profiler
        or by a pipeline that recorded onto ``cache.event_profiler``.
        Otherwise a snapshot of this AR step's per-stage timings (ms) and
        GPU memory (GiB):
        ``{<stage>_ms, total_ms, total_ms_wo_finalize, mem_alloc_gib,
        mem_reserved_gib, mem_peak_gib}``. The same numbers are also
        logged via ``logger.info``.
    """
    assert cache.autoregressive_index == autoregressive_index, (
        f"autoregressive_index mismatch: generate() ran with "
        f"{cache.autoregressive_index} but finalize() was called with "
        f"{autoregressive_index}."
    )
    assert cache.final_state is not None, (
        "finalize() called before generate() — no FinalState on the cache."
    )
    profiler = get_inference_profiler()
    with profiler.stage_scope(cache), profiler.range("pipeline.finalize"):
        self.diffusion_model.finalize(final_state=cache.final_state)
    if cache.event_profiler is not None:
        # A pipeline that overrides generate() times its own stages onto
        # the cache. Those are the ones to report.
        cache.event_profiler.record("finalize")
        stats_ms = cache.event_profiler.sync_and_summarize()
    else:
        with profiler.stage_scope(cache):
            stats_ms = profiler.collect_stage_ms()
    if not stats_ms:
        return None

    stats = format_result_as_ms(
        stats_ms, collect_totals=True, collect_vram_info=True
    )

    # Custom logging
    total_ms = stats["total_ms"]
    total_ms_wo_finalize = stats["total_ms_wo_finalize"]
    stages_str = " ".join(f"{stage} {ms:.3f} ms" for stage, ms in stats_ms.items())
    mem_str = ""
    if "mem_alloc_gib" in stats:
        mem_str = (
            f" | GPU mem alloc {stats['mem_alloc_gib']:.3f} GiB "
            f"reserved {stats['mem_reserved_gib']:.3f} GiB "
            f"peak {stats['mem_peak_gib']:.3f} GiB"
        )
    logger.info(
        f"AR {autoregressive_index} {stages_str} | "
        f"total(w/o finalize) {total_ms_wo_finalize:.3f} ms "
        f"total {total_ms:.3f} ms"
        f"{mem_str}"
    )
    # End of custom logging

    return stats

def generate(
    self,
    autoregressive_index: int,
    cache: TransformerCacheT,
    input: Any = None,
) -> tuple[Tensor, "DiffusionModel.FinalState[TransformerCacheT]"]:
    """Run the denoising loop for one AR step.

    Args:
        autoregressive_index: AR step index.
        cache: Long-lived AR cache, mutated in place.
        input: Optional per-AR-step encoder output. Patchified here and
            forwarded to ``predict_flow`` / ``postprocess_clean_latent``,
            then stashed on the returned ``FinalState`` for ``finalize``.

    Returns:
        ``(clean_latent, final_state)``. ``clean_latent`` is unpatchified;
        ``final_state`` should be passed to ``finalize``.
    """
    if input is not None:
        input = self.transformer.patchify_and_maybe_split_cp(input)
    cache.start(autoregressive_index)

    if self.config.noise_in_unpatchified_shape:
        # Draw noise at the unpatchified shape so the RNG sequence matches
        # implementations that sample before patchify; ``self.latent_shape``
        # is patchified, so go through ``unpatchify_and_maybe_gather_cp``
        # first to recover the pixel-side latent shape.
        dummy_latent = torch.empty(
            self.latent_shape, device=self.device, dtype=self.dtype
        )
        dummy_latent = self.transformer.unpatchify_and_maybe_gather_cp(dummy_latent)
        initial_noise = self.transformer.initial_noise(
            latent_shape=tuple(dummy_latent.shape),
            rng=self.rng,
            cache=cache,
            input=input,
        )

    else:
        initial_noise = self.transformer.initial_noise(
            latent_shape=self.latent_shape,
            rng=self.rng,
            cache=cache,
            input=input,
        )

    def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor:
        # Round-trip patchify/unpatchify around the network when the
        # scheduler is operating on unpatchified latents; the transformer
        # itself always consumes/emits patchified shapes.
        if self.config.noise_in_unpatchified_shape:
            noisy_latent = self.transformer.patchify_and_maybe_split_cp(
                noisy_latent
            )

        output = self.transformer.predict_flow(
            noisy_latent=noisy_latent,
            timestep=timestep,
            cache=cache,
            input=input,
        )

        if self.config.noise_in_unpatchified_shape:
            output = self.transformer.unpatchify_and_maybe_gather_cp(output)
        return output

    clean_latent = self.scheduler.sample(
        initial_noise=initial_noise,
        predict_flow=predict_flow,
        rng=self.rng,
    )

    if self.config.noise_in_unpatchified_shape:
        # Scheduler emitted an unpatchified latent; patchify before
        # ``postprocess_clean_latent`` and the cache-update path.
        clean_latent = self.transformer.patchify_and_maybe_split_cp(clean_latent)

    clean_latent = self.transformer.postprocess_clean_latent(
        clean_latent=clean_latent,
        cache=cache,
        input=input,
    )

    # Retain the clean state for finalization. A transformer may already
    # have refreshed specialized cache state in ``postprocess_clean_latent``.
    final_state = DiffusionModel.FinalState(
        clean_latent=clean_latent,
        autoregressive_index=autoregressive_index,
        cache=cache,
        input=input,
    )

    clean_latent = self.transformer.unpatchify_and_maybe_gather_cp(clean_latent)
    return clean_latent, final_state

def finalize(
    self,
    final_state: "DiffusionModel.FinalState[TransformerCacheT]",
) -> None:
    """Advance the AR cache using the clean latent from ``generate``.

    Re-noises the clean latent to ``config.context_noise`` and runs the
    transformer's ``finalize_kv_cache`` (one forward for vanilla
    transformers, multiple for dual-network DiTs).

    ``context_noise == 0`` skips ``add_noise`` (sigma=0 is identity) and
    feeds the clean latent directly. This also dodges the requirement
    for schedulers to support a ``t=0`` lookup (UniPC's inference
    schedule has no ``t=0`` entry).
    """
    context_noise = self.config.context_noise
    timestep = torch.tensor(context_noise, device=self.device, dtype=self.dtype)
    if context_noise > 0.0:
        clean_latent = final_state.clean_latent
        # Round-trip patchify/unpatchify around ``add_noise`` so the
        # scheduler sees the same shape as in ``generate``; the cache
        # update below always consumes a patchified latent.
        if self.config.noise_in_unpatchified_shape:
            clean_latent = self.transformer.unpatchify_and_maybe_gather_cp(
                final_state.clean_latent
            )
        noisy_latent = self.scheduler.add_noise(
            clean_input=clean_latent,
            timestep=timestep,
            rng=self.rng,
        )
        if self.config.noise_in_unpatchified_shape:
            noisy_latent = self.transformer.patchify_and_maybe_split_cp(
                noisy_latent
            )
    else:
        noisy_latent = final_state.clean_latent
    self.transformer.finalize_kv_cache(
        noisy_latent=noisy_latent,
        timestep=timestep,
        cache=final_state.cache,
        input=final_state.input,
    )
    final_state.cache.finalize(final_state.autoregressive_index)

@abstractmethod
def predict_flow(
    self,
    noisy_latent: Tensor,
    timestep: Tensor,
    cache: TransformerCacheT,
    input: Any = None,
) -> Tensor:
    """Predict the flow at ``timestep``.

    Args:
        noisy_latent: Patchified noisy latent for this denoising step.
        timestep: Scalar timestep tensor.
        cache: Per-rollout AR cache.
        input: Patchified encoder output for this AR step, or ``None``
            when the pipeline has no encoder. Subclasses should narrow
            the type to their encoder's output type.

    Returns:
        Predicted flow tensor with the same shape as ``noisy_latent``.
    """

def finalize_kv_cache(
    self,
    noisy_latent: Tensor,
    timestep: Tensor,
    cache: TransformerCacheT,
    input: Any = None,
) -> None:
    """Advance the AR cache so it is ready for the next AR step.

    Called by ``DiffusionModel.finalize`` after the denoising loop; the
    flow is discarded — only the cache side effect matters. Default
    runs a single ``predict_flow`` forward. Override for transformers
    with multiple parallel networks that must stay in lock-step.

    Args:
        noisy_latent: Patchified latent at the AR-step's context noise
            (or the clean latent when ``context_noise == 0``).
        timestep: 0-d context-noise timestep tensor.
        cache: Per-rollout AR cache.
        input: Same patchified encoder output passed to ``predict_flow``.
    """
    _ = self.predict_flow(noisy_latent, timestep, cache, input)

def initialize_autoregressive_cache(self, **context: Any) -> TransformerCacheT:
    """Build a fresh AR cache for a new rollout.

    Default returns an empty cache, correct for non-AR transformers.
    Subclasses with custom cache types must override this and may
    declare typed per-rollout context (e.g. ``text_embeddings``).

    Args:
        context: Per-rollout state forwarded as keyword arguments.

    Returns:
        Fresh AR cache.
    """
    return cast("TransformerCacheT", TransformerAutoregressiveCache())

def initial_noise(
    self,
    *,
    latent_shape: tuple[int, ...],
    rng: torch.Generator | None,
    cache: TransformerCacheT,
    input: Any = None,
) -> Tensor:
    """Draw the starting latent for one denoising loop.

    Default is Gaussian noise in ``latent_shape``. Model integrations can
    override this to seed from encoder outputs, pin I2V frames, or attach
    per-step state to the cache while still using ``DiffusionModel``.

    Args:
        latent_shape: Shape requested by the diffusion model.
        rng: Per-model generator on ``self.device``, or ``None``.
        cache: Per-rollout AR cache.
        input: Same patchified encoder output passed to ``predict_flow``.

    Returns:
        Initial latent tensor for scheduler sampling.
    """
    del cache, input
    return torch.randn(
        latent_shape,
        device=self.device,
        dtype=self.dtype,
        generator=rng,
    )

def postprocess_clean_latent(
    self,
    clean_latent: Tensor,
    cache: TransformerCacheT,
    input: Any = None,
) -> Tensor:
    """Optional postprocessing hook for the predicted clean latent.

    Default is identity. Override to clamp or re-inject regions whose
    clean value is known a priori (e.g. I2V first-frame pinning).
    Called at the end of ``DiffusionModel.generate``.

    Args:
        clean_latent: Patchified ``x0`` from the denoising loop.
        cache: Per-rollout AR cache.
        input: Same patchified encoder output passed to ``predict_flow``.

    Returns:
        Postprocessed clean latent with the same shape.
    """
    return clean_latent

@abstractmethod
def patchify_and_maybe_split_cp(self, x: Any) -> Any:
    """Patchify and (optionally) CP-split a noisy latent or encoder payload.

    Tensors patchify and split. Structured payloads (e.g. an image-control
    struct with ``latent`` + ``mask``) patchify each tensor field and
    return the same struct type. Implement as identity when neither
    token packing nor CP sharding applies. Output preserves the input
    Python type — only shapes change.
    """

@abstractmethod
def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor:
    """Inverse of ``patchify_and_maybe_split_cp`` for the network output."""

def start(self, autoregressive_index: int) -> None:
    """Mark the start of an AR step. Default is a no-op.

    Args:
        autoregressive_index: Index of the AR step being started.
    """

def finalize(self, autoregressive_index: int) -> None:
    """Finalize bookkeeping after use at this AR step. Default is a no-op.

    Args:
        autoregressive_index: Index of the AR step just finalized.
    """

@abstractmethod
def sample(
    self,
    initial_noise: Tensor,
    predict_flow: FlowPredictor,
    rng: torch.Generator | None = None,
) -> Tensor:
    """Run the full denoising loop and return the clean latent.

    Schedulers are shape-agnostic: every internal op broadcasts against
    per-step scalar sigmas. In practice ``initial_noise`` is a video
    latent ``[B, C, T, H, W]``, conventionally treated as a sample at
    ``sigma=1``.

    Args:
        initial_noise: Gaussian noise on the caller's device/dtype.
        predict_flow: Per-step closure invoked
            ``num_inference_steps`` times.
        rng: Generator on the same device. Used by self-forcing renoise
            loops; pure ODE solvers ignore it.

    Returns:
        Clean latent with the same shape, device, and dtype as
        ``initial_noise``.
    """

@abstractmethod
def add_noise(
    self,
    clean_input: Tensor,
    timestep: Tensor,
    rng: torch.Generator | None = None,
) -> Tensor:
    """Apply the forward corruption ``x_t = (1 - sigma(t)) * x_0 + sigma(t) * eps``.

    Timestep value semantics are scheduler-specific: all schedulers snap
    to the nearest entry of their inference schedule.
    """

def __call__(self, noisy_latent: Tensor, timestep: Tensor) -> Tensor:
    """Predict the flow at ``timestep``."""
    ...

def sample(
    self,
    initial_noise: Tensor,
    predict_flow: FlowPredictor,
    rng: torch.Generator | None = None,
) -> Tensor:
    """Run the self-forcing flow-match denoising loop.

    Iteration 0 trusts ``initial_noise`` as the ``sigma=1`` sample;
    later iterations re-noise the previous ``x0`` estimate to the new
    sigma before the network forward. By default the result is cast back
    to ``initial_noise.dtype`` after every step; reference sampling casts
    each ``x0`` to the network flow dtype.
    """
    input_dtype = initial_noise.dtype
    reference_sampling = self.config.reference_sampling
    sigmas = self.denoising_sigmas
    timesteps = self.denoising_step_list

    noisy = initial_noise
    clean: Tensor | None = None
    for i in tqdm(
        range(timesteps.shape[0]),
        disable=not self.config.enable_tqdm,
        desc="FlowMatchScheduler",
    ):
        sigma = sigmas[i]
        # Schedule buffers are pinned to fp32 (to preserve integer
        # timestep values under a stray `module.to(bf16)`), but the
        # By default the network expects timesteps in the input dtype.
        # Reference sampling preserves configured integer values rather
        # than rounding them through bf16.
        timestep = (
            timesteps[i]
            if reference_sampling
            else timesteps[i].to(dtype=input_dtype)
        )
        if i > 0:
            assert clean is not None
            if reference_sampling:
                noise = torch.empty_like(clean).normal_(generator=rng)
                state_dtype = clean.dtype
            else:
                noise = torch.empty_like(noisy).normal_(generator=rng)
                state_dtype = input_dtype
            noisy = ((1.0 - sigma) * clean + sigma * noise).to(state_dtype)
        flow = predict_flow(noisy, timestep)
        if reference_sampling:
            clean = (noisy - sigma * flow).to(flow.dtype)
        else:
            clean = noisy - sigma * flow
    assert clean is not None, "denoising_step_list is empty"
    return clean if reference_sampling else clean.to(input_dtype)

def add_noise(
    self,
    clean_input: Tensor,
    timestep: Tensor,
    rng: torch.Generator | None = None,
) -> Tensor:
    """Apply the forward corruption at an arbitrary timestep.

    Snaps ``timestep`` to the nearest entry of the warped training table
    and uses it as sigma in the standard lerp.
    """
    assert timestep.shape == (), f"expected scalar timestep, got {timestep.shape}"
    full_t = self._full_timesteps
    idx = torch.argmin((full_t - timestep.to(full_t.dtype)).abs()).reshape(1)
    sigma = self._full_sigmas.index_select(0, idx).reshape(())
    noise = torch.empty_like(clean_input).normal_(generator=rng)
    return ((1.0 - sigma) * clean_input + sigma * noise).to(clean_input.dtype)

def sample(
    self,
    initial_noise: Tensor,
    predict_flow: FlowPredictor,
    rng: torch.Generator | None = None,
) -> Tensor:
    """Run the order-2 UniPC predictor-corrector denoising loop.

    Each iteration: network → flow → ``x0`` → corrector (skipped at
    step 0) → predictor. All per-step coefficients are pre-baked at
    construction; the loop is pure tensor ops. Internal arithmetic is
    fp32; the result is cast back to ``initial_noise.dtype``. ``rng``
    is unused (deterministic ODE) but accepted for interface conformance.
    """
    input_dtype = initial_noise.dtype
    N = self.timesteps.shape[0]

    sample = initial_noise
    m_prev: Tensor | None = None
    m_prev_prev: Tensor | None = None
    last_sample: Tensor | None = None

    for i in tqdm(
        range(N),
        disable=not self.config.enable_tqdm,
        desc="FlowMatchUniPCScheduler",
    ):
        # Schedule buffers are pinned to fp32 (to preserve integer
        # timestep values under a stray `module.to(bf16)`), but the
        # network expects timesteps in the input dtype so that
        # downstream modulation / Linear layers stay consistent.
        timestep = self.timesteps[i].to(dtype=input_dtype)

        # Network forward (heavy compute -- everything else here is
        # ~free relative to this).
        flow = predict_flow(sample, timestep)

        # Convert to x0 estimate (predict_x0 + flow_prediction):
        #   m_curr = sample - sigma_i * flow
        # Promote to fp32 to match upstream's
        # ``model_output = model_output.to(dtype=torch.float32)``
        # before convert_model_output.
        m_curr = sample.to(torch.float32) - self.sigmas[i] * flow.to(torch.float32)

        # Corrector (skip on first step).
        #
        # last_sample is the sample saved before the previous
        # predictor (matches upstream ``self.last_sample``).
        # m_prev is model_outputs[-1] (previous step's m_curr),
        # m_prev_prev is model_outputs[-2] (2 steps back).
        # b_corr_dprev[1] = 0 (order-1 corrector) so the
        # m_prev_prev term vanishes at i=1; we alias it to m_prev
        # there to skip the zero-tensor allocation.
        if i >= 1:
            assert last_sample is not None and m_prev is not None
            m_pp = m_prev_prev if m_prev_prev is not None else m_prev
            corrected = (
                self.a_corr[i] * last_sample.to(torch.float32)
                + self.b_corr_m0[i] * m_prev
                + self.b_corr_dprev[i] * (m_pp - m_prev)
                + self.b_corr_dt[i] * (m_curr - m_prev)
            )
            sample = corrected.to(input_dtype)

        # Save sample BEFORE the predictor so the next iteration's
        # corrector can use it as "last_sample".
        last_sample = sample

        # Predictor advances sample to next sigma:
        #   x_{i+1} = a_pred[i] * x_i + b_pred_m0[i] * m0
        #            + b_pred_dprev[i] * (m_prev - m0)
        # b_pred_dprev[0] = 0 (order-1 warmup) and
        # b_pred_dprev[-1] = 0 (order-1 final) so the same line
        # serves both order branches; alias m_prev to m_curr at
        # the warmup step to skip a zero-tensor allocation.
        m_p = m_prev if m_prev is not None else m_curr
        predicted = (
            self.a_pred[i] * sample.to(torch.float32)
            + self.b_pred_m0[i] * m_curr
            + self.b_pred_dprev[i] * (m_p - m_curr)
        )
        sample = predicted.to(input_dtype)

        # Roll the model-output history.
        m_prev_prev = m_prev
        m_prev = m_curr

    return sample