TestingRef2va / pk_workflow.py
dagloop5's picture
Upload pk_workflow.py
0925136 verified
Raw
History Blame Contribute Delete
24.4 kB
"""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