"""The three things that make `Plaguekind/Minimax-H3` a *workflow* rather than just MiniMax-H3.

`Plaguekind/Minimax-H3` ships no weights: it is a ComfyUI graph
(`PlagueKind-MinimaxH3-V1.5.json`) over `Comfy-Org/MiniMax-H3`, and everything it contributes is in the
sampling and the post chain. Read off the graph, that is:

| ComfyUI node | widget | here |
|---|---|---|
| `KSamplerSelect` | `euler` | MiniMax-H3's only sampler; the checkpoint is CFG-distilled, one forward per step |
| `BasicScheduler` | `linear_quadratic`, 15 steps, denoise 1.0 | `linear_quadratic_sigmas` |
| `ImageSharpenKJ` | `rcas`, 0.3 | `rcas` |
| `FrameInterpolate` + `FrameInterpolationModelLoader` | `film_net_fp16.safetensors`, multiplier 2 | `interpolate` |
| `CreateVideo` | fps `24 * 2` | 48 fps out |
| `RTXVideoSuperResolution` | 2x, `ULTRA` | **not portable** — NVIDIA NGX, Windows/RTX driver only |

The sigma schedule is the one that changes the pixels most, and the one that is easy to get subtly wrong.
"""

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 _BatchedBrownianTree:
    """Minimal port of k-diffusion's `BatchedBrownianTree` (single-seed case only — MiniMax-H3 requests run at
    batch size 1). Wraps `torchsde.BrownianTree` so consecutive noise draws at adjacent sigma pairs are
    correlated through a shared stochastic path, as `dpmpp_2m_sde`/`dpmpp_3m_sde` require — independent
    per-step Gaussian noise (as used for `euler_ancestral`/`er_sde` above) is a materially different sampler.
    """

    def __init__(self, x, t0, t1, seed, cpu=False):
        import torchsde

        self.cpu_tree = cpu
        if t0 > t1:
            t0, t1, self.sign = t1, t0, -1
        else:
            self.sign = 1
        w0 = torch.zeros_like(x)
        if self.cpu_tree:
            t0, w0, t1 = t0.detach().cpu(), w0.detach().cpu(), t1.detach().cpu()
        self.tree = torchsde.BrownianTree(t0, w0, t1, entropy=seed)

    def __call__(self, t0, t1):
        if t0 > t1:
            t0, t1, sign = t1, t0, -1
        else:
            sign = 1
        device, dtype = t0.device, t0.dtype
        if self.cpu_tree:
            t0, t1 = t0.detach().cpu().float(), t1.detach().cpu().float()
        return self.tree(t0, t1).to(device=device, dtype=dtype) * (self.sign * sign)


class _BrownianTreeNoiseSampler:
    """Port of k-diffusion's `BrownianTreeNoiseSampler`. `cpu=False` matches the `*_gpu` sampler names — noise
    is generated directly on the accelerator rather than the CPU-tree variant the non-`_gpu` names use."""

    def __init__(self, x, sigma_min, sigma_max, seed, cpu=False):
        self.tree = _BatchedBrownianTree(x, torch.as_tensor(sigma_min), torch.as_tensor(sigma_max), seed, cpu=cpu)

    def __call__(self, sigma, sigma_next):
        t0, t1 = torch.as_tensor(sigma), torch.as_tensor(sigma_next)
        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()

def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
    """Ports k-diffusion's `sample_dpmpp_2m_sde` (`solver_type='midpoint'`, matching ComfyUI's `dpmpp_2m_sde_gpu`
    — the `_heun` variant is a different `solver_type` and is not ported here) onto one `MiniMaxH3Scheduler.step()`
    call. Single model evaluation per step; second-order accuracy comes from the previous step's denoised
    estimate. History and the Brownian-tree noise sampler live on the scheduler instance, reset each request.
    """
    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
    # Cached once and reused every call — `torchsde.BrownianTree` caches its internal tree keyed to the exact
    # float value it was first queried with, and re-deriving "the same" sigma via a fresh `.to()` cast on a
    # later call can land a few ULPs away from what the tree remembers, which it treats as an ordering error.
    if scheduler._dpmpp_sde_sigmas is None:
        scheduler._dpmpp_sde_sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
    sigmas = scheduler._dpmpp_sde_sigmas
    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:
        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
        shift = float(scheduler.shift)
        sigma = shift * base / (1 + (shift - 1) * base)

    if scheduler._dpmpp_sde_noise_sampler is None:
        # `cpu=True` runs the Brownian-bridge recursion on CPU rather than the GPU — negligible cost next to
        # the transformer forward pass, but noticeably more numerically stable than `cpu=False`, which is what
        # ComfyUI's own non-`_gpu`-suffixed variants default to for exactly this reason.
        scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
        )

    def half_log_snr(s):
        return torch.log((1 - s) / s)

    if sigma_next == 0:
        prev_sample = denoised
    else:
        lambda_s, lambda_t = half_log_snr(sigma), half_log_snr(sigma_next)
        h = lambda_t - lambda_s
        h_eta = h * (eta + 1)
        alpha_next = 1 - sigma_next

        prev_sample = (sigma_next / sigma) * (-h * eta).exp() * x + alpha_next * (-h_eta).expm1().neg() * denoised

        old_denoised = scheduler._dpmpp_sde_old_denoised
        h_last = scheduler._dpmpp_sde_h_last
        if old_denoised is not None:
            r = h_last / h
            prev_sample = prev_sample + 0.5 * alpha_next * (-h_eta).expm1().neg() * (1 / r) * (denoised - old_denoised)

        if eta > 0 and s_noise > 0:
            noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
            prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise

        scheduler._dpmpp_sde_h_last = h

    scheduler._dpmpp_sde_old_denoised = denoised
    prev_sample = prev_sample.to(dtype=sample.dtype)
    scheduler._step_index += 1
    return prev_sample


def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
    """Ports k-diffusion's `sample_dpmpp_3m_sde` (matching ComfyUI's `dpmpp_3m_sde_gpu`) onto one
    `MiniMaxH3Scheduler.step()` call. Single model evaluation per step; third-order accuracy (once two prior
    steps exist) comes from history carried on the scheduler instance, plus the same Brownian-tree noise as
    `_dpmpp_2m_sde_step`.
    """
    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
    # Cached once and reused every call — `torchsde.BrownianTree` caches its internal tree keyed to the exact
    # float value it was first queried with, and re-deriving "the same" sigma via a fresh `.to()` cast on a
    # later call can land a few ULPs away from what the tree remembers, which it treats as an ordering error.
    if scheduler._dpmpp_sde_sigmas is None:
        scheduler._dpmpp_sde_sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
    sigmas = scheduler._dpmpp_sde_sigmas
    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:
        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
        shift = float(scheduler.shift)
        sigma = shift * base / (1 + (shift - 1) * base)

    if scheduler._dpmpp_sde_noise_sampler is None:
        # `cpu=True` runs the Brownian-bridge recursion on CPU rather than the GPU — negligible cost next to
        # the transformer forward pass, but noticeably more numerically stable than `cpu=False`, which is what
        # ComfyUI's own non-`_gpu`-suffixed variants default to for exactly this reason.
        scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
        )

    def half_log_snr(s):
        return torch.log((1 - s) / s)

    if sigma_next == 0:
        prev_sample = denoised
    else:
        lambda_s, lambda_t = half_log_snr(sigma), half_log_snr(sigma_next)
        h = lambda_t - lambda_s
        h_eta = h * (eta + 1)
        alpha_next = 1 - sigma_next

        prev_sample = (sigma_next / sigma) * (-h * eta).exp() * x + alpha_next * (-h_eta).expm1().neg() * denoised

        denoised_1 = scheduler._dpmpp_sde_old_denoised
        denoised_2 = scheduler._dpmpp_sde_old_denoised_2
        h_1 = scheduler._dpmpp_sde_h_last
        h_2 = scheduler._dpmpp_sde_h_last_2

        if h_2 is not None:
            r0, r1 = h_1 / h, h_2 / h
            d1_0 = (denoised - denoised_1) / r0
            d1_1 = (denoised_1 - denoised_2) / r1
            d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
            d2 = (d1_0 - d1_1) / (r0 + r1)
            phi_2 = h_eta.neg().expm1() / h_eta + 1
            phi_3 = phi_2 / h_eta - 0.5
            prev_sample = prev_sample + (alpha_next * phi_2) * d1 - (alpha_next * phi_3) * d2
        elif h_1 is not None:
            r = h_1 / h
            d = (denoised - denoised_1) / r
            phi_2 = h_eta.neg().expm1() / h_eta + 1
            prev_sample = prev_sample + (alpha_next * phi_2) * d

        if eta > 0 and s_noise > 0:
            noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
            prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise

        scheduler._dpmpp_sde_h_last_2 = h_1
        scheduler._dpmpp_sde_h_last = h

    scheduler._dpmpp_sde_old_denoised_2 = scheduler._dpmpp_sde_old_denoised
    scheduler._dpmpp_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
        elif self.sampler_name in ("dpmpp_2m_sde_gpu", "dpmpp_3m_sde_gpu"):
            step_fn = _dpmpp_2m_sde_step if self.sampler_name == "dpmpp_2m_sde_gpu" else _dpmpp_3m_sde_step
            for offset, attr_name in enumerate(self.attr_names):
                scheduler = getattr(self.pipe, attr_name)
                scheduler._dpmpp_sde_old_denoised = None
                scheduler._dpmpp_sde_old_denoised_2 = None
                scheduler._dpmpp_sde_h_last = None
                scheduler._dpmpp_sde_h_last_2 = None
                scheduler._dpmpp_sde_noise_sampler = None
                scheduler._dpmpp_sde_sigmas = None
                scheduler._dpmpp_sde_seed = self.seed + offset

                def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _f=step_fn, **_kwargs):
                    return (_f(_s, 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