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