File size: 24,430 Bytes
0925136
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
"""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