Spaces:
Running on Zero
Running on Zero
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
|