392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""Shared machinery for the dual-schedule sampler (MiniMax H3).
|
|
|
|
H3 denoises a video stream and an audio stream in one DiT forward, and the two were trained
|
|
on different flow shifts (12.0 video / 3.0 audio). Core keeps a single SIGMAS input by
|
|
carrying the audio latent onto the video schedule, which locks the audio schedule to a
|
|
re-shift of the video one. This sampler takes both schedules and steps each stream on its
|
|
own, without patching core:
|
|
|
|
- the carry is switched off by shadowing MiniMaxH3.audio_scale with 1.0 for the run, so the
|
|
sampler holds the real audio latent at its own sigma
|
|
- the audio noise level reaches the DiT through the flow shifts it already reads from
|
|
transformer_options: the video shift is pinned to 1.0 (identity) and the audio shift
|
|
solved so the model's own mapping lands on the step's audio sigma
|
|
- the audio stream's x0 is re-derived from the returned velocity at the audio sigma, since
|
|
calculate_denoised applies the single timestep sigma to the whole pack
|
|
|
|
Both schedules are used exactly as given. An all-zero schedule freezes its stream: the
|
|
latent is left alone and the DiT is told that stream is clean, which is how one stream is
|
|
refined while the other only conditions it.
|
|
"""
|
|
|
|
import copy
|
|
import math
|
|
|
|
from comfy.k_diffusion import sampling as k_diffusion_sampling
|
|
import comfy.model_sampling
|
|
import comfy.samplers
|
|
from comfy_api.latest import io
|
|
import torch
|
|
from tqdm.auto import trange
|
|
|
|
# A socket of its own, so a stock SAMPLER cannot be wired into a dual sampling node.
|
|
DualSamplerType = io.Custom("DUAL_SAMPLER")
|
|
|
|
|
|
def audio_shift_for(sigma_v, sigma_a):
|
|
"""The to_shift that makes time_shift_sigma(sigma_v, 1.0, to_shift) == sigma_a.
|
|
|
|
With from_shift 1.0 the model's inversion is the identity, so its mapping reduces to
|
|
sigma_a = to*sigma_v / (1 + (to-1)*sigma_v), which inverts in closed form. The mapping
|
|
fixes sigma = 1, which check_schedules rejects, so the guard here is only numerical.
|
|
"""
|
|
if sigma_v >= 1.0 or sigma_a >= 1.0:
|
|
return 1.0
|
|
return sigma_a * (1.0 - sigma_v) / (sigma_v * (1.0 - sigma_a))
|
|
|
|
|
|
def is_frozen(sigmas):
|
|
"""An all-zero schedule means: leave this stream alone, it is already x0."""
|
|
return bool((sigmas == 0).all())
|
|
|
|
|
|
def check_schedules(video_sigmas, audio_sigmas):
|
|
"""Reject the two schedule pairings that cannot be sampled. Nothing is rewritten."""
|
|
if video_sigmas.shape[-1] != audio_sigmas.shape[-1]:
|
|
raise ValueError(
|
|
"video_sigmas and audio_sigmas must have the same number of steps, got {} and {}.".format(
|
|
video_sigmas.shape[-1], audio_sigmas.shape[-1]
|
|
)
|
|
)
|
|
if video_sigmas.shape[-1] == 0:
|
|
return
|
|
# The model maps between the two streams' noise levels with a transform that fixes 1.0,
|
|
# so at 1.0 both streams are pinned together and the other one cannot be expressed.
|
|
first_v, first_a = float(video_sigmas[0]), float(audio_sigmas[0])
|
|
if first_v >= 1.0 and first_a < 1.0:
|
|
raise ValueError(
|
|
"Video sigma starts at 1.0 while the audio sigma is less than 1 ({:g}). Use 0.9999 for the first video sigma.".format(
|
|
first_a
|
|
)
|
|
)
|
|
if first_a >= 1.0 and first_v < 1.0:
|
|
raise ValueError(
|
|
"Audio sigma starts at 1.0 while the video sigma is less than 1 ({:g}). Use 0.9999 for the first audio sigma.".format(
|
|
first_v
|
|
)
|
|
)
|
|
|
|
|
|
def ancestral_stream(x, sigmas, o):
|
|
"""Ancestral Euler over one stream, mirroring core's sample_euler_ancestral_RF.
|
|
|
|
A coroutine: it yields (step, sigma, x) to ask for this stream's x0 estimate at that
|
|
noise level and receives it back. The driver pairs both streams' requests into a single
|
|
DiT forward, which is what lets the two run on independent schedules. eta 0 makes it
|
|
plain Euler.
|
|
"""
|
|
for i in range(len(sigmas) - 1):
|
|
denoised = yield (i, sigmas[i], x)
|
|
if sigmas[i + 1] == 0:
|
|
x = denoised
|
|
continue
|
|
downstep_ratio = 1 + (sigmas[i + 1] / sigmas[i] - 1) * o.eta
|
|
sigma_down = sigmas[i + 1] * downstep_ratio
|
|
alpha_ip1 = 1 - sigmas[i + 1]
|
|
alpha_down = 1 - sigma_down
|
|
renoise_coeff = (
|
|
sigmas[i + 1] ** 2 - sigma_down**2 * alpha_ip1**2 / alpha_down**2
|
|
) ** 0.5
|
|
sigma_down_i_ratio = sigma_down / sigmas[i]
|
|
x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * denoised
|
|
if o.eta > 0:
|
|
x = (alpha_ip1 / alpha_down) * x + o.noise(
|
|
sigmas[i], sigmas[i + 1]
|
|
) * o.s_noise * renoise_coeff
|
|
return x
|
|
|
|
|
|
class StreamOptions:
|
|
"""Per-stream knobs and noise source handed to the sampler coroutine."""
|
|
|
|
def __init__(self, eta, s_noise, noise):
|
|
self.eta, self.s_noise, self.noise = eta, s_noise, noise
|
|
|
|
|
|
class DualSampler(comfy.samplers.Sampler):
|
|
"""Steps a packed [video, audio] latent on two independent sigma schedules.
|
|
|
|
The video schedule arrives through the normal `sigmas` argument; the audio schedule is
|
|
attached by DualSamplerCustomAdvanced via with_audio_sigmas().
|
|
"""
|
|
|
|
def __init__(
|
|
self, eta_video=1.0, eta_audio=1.0, s_noise_video=1.0, s_noise_audio=1.0
|
|
):
|
|
self.eta_video = eta_video
|
|
self.eta_audio = eta_audio
|
|
self.s_noise_video = s_noise_video
|
|
self.s_noise_audio = s_noise_audio
|
|
self.audio_sigmas = None
|
|
|
|
def with_audio_sigmas(self, audio_sigmas):
|
|
sampler = copy.copy(self)
|
|
sampler.audio_sigmas = audio_sigmas
|
|
return sampler
|
|
|
|
def stream_options(self, x, split, video, seed, model_sampling):
|
|
"""Noise source for one stream.
|
|
|
|
The sampler is built on the *packed* shape and sliced, so with equal schedules the
|
|
two streams draw exactly the tensor core's single sampler would have drawn.
|
|
"""
|
|
region = slice(None, split) if video else slice(split, None)
|
|
draw = k_diffusion_sampling.default_noise_sampler(x, seed=seed)
|
|
noise_scale = getattr(model_sampling, "noise_scale", 1.0)
|
|
return StreamOptions(
|
|
eta=self.eta_video if video else self.eta_audio,
|
|
s_noise=(self.s_noise_video if video else self.s_noise_audio)
|
|
* noise_scale,
|
|
noise=lambda a, b: draw(a, b)[..., region],
|
|
)
|
|
|
|
def sample(
|
|
self,
|
|
model_wrap,
|
|
sigmas,
|
|
extra_args,
|
|
callback,
|
|
noise,
|
|
latent_image=None,
|
|
denoise_mask=None,
|
|
disable_pbar=False,
|
|
):
|
|
if self.audio_sigmas is None:
|
|
raise ValueError(
|
|
"A dual sampler needs both schedules: connect it to DualSamplerCustomAdvanced, not SamplerCustomAdvanced."
|
|
)
|
|
|
|
inner_model = model_wrap.inner_model
|
|
model_sampling = inner_model.model_sampling
|
|
if not isinstance(model_sampling, comfy.model_sampling.CONST):
|
|
raise ValueError(
|
|
"The dual sampler only supports rectified flow (CONST) models such as MiniMax H3."
|
|
)
|
|
# The shifts on the guider's model are irrelevant here - this node overrides both per
|
|
# step and drives the streams from the input sigmas. The multiplier is not: it is the
|
|
# timestep scale the model is called on, and the H3 DiT reads sigma back as
|
|
# timestep/1000. ModelSamplingAuraFlow patches at multiplier 1.0, which the stock path
|
|
# rejects for other reasons but this one would otherwise run 1000x off in silence.
|
|
multiplier = getattr(model_sampling, "multiplier", None)
|
|
if multiplier != 1000:
|
|
raise ValueError(
|
|
"The guider's model is patched at timestep multiplier {} instead of 1000, so the DiT would read every sigma too small. Use ModelSamplingMiniMaxH3 or ModelSamplingSD3 on the guider and keep AuraFlow on a separate branch feeding BasicScheduler.".format(
|
|
multiplier
|
|
)
|
|
)
|
|
|
|
latent_shapes = inner_model.latent_shapes
|
|
if latent_shapes is None or len(latent_shapes) < 2:
|
|
raise ValueError(
|
|
"DualSamplerCustomAdvanced needs a packed audio-video latent (two streams), got one stream."
|
|
)
|
|
|
|
# pack_latents lays the streams out flat and in order: video first, then audio
|
|
split = math.prod(latent_shapes[0][1:])
|
|
|
|
video_sigmas = sigmas
|
|
audio_sigmas = self.audio_sigmas.to(
|
|
device=video_sigmas.device, dtype=video_sigmas.dtype
|
|
)
|
|
|
|
model_options = extra_args.get("model_options", {})
|
|
# a per-sample deep clone, so per-step writes stay local to this run
|
|
transformer_options = model_options.setdefault(
|
|
"transformer_options", {}
|
|
)
|
|
transformer_options["sample_sigmas_audio"] = audio_sigmas
|
|
|
|
seed = extra_args.get("seed", None)
|
|
|
|
check_schedules(video_sigmas, audio_sigmas)
|
|
total_steps = len(video_sigmas) - 1
|
|
|
|
# an all-zero schedule freezes its stream: the latent is already x0, so it is never
|
|
# stepped and the DiT is told t = 1 (clean) for it every pass
|
|
frozen_v, frozen_a = is_frozen(video_sigmas), is_frozen(audio_sigmas)
|
|
if frozen_v and frozen_a:
|
|
raise ValueError(
|
|
"both schedules are all zeros, so there is nothing to sample."
|
|
)
|
|
|
|
# each stream is brought up to its own first sigma, so the two can start anywhere
|
|
x = torch.empty_like(noise)
|
|
x[..., :split] = model_sampling.noise_scaling(
|
|
video_sigmas[0],
|
|
noise[..., :split],
|
|
latent_image[..., :split],
|
|
self.max_denoise(model_wrap, video_sigmas),
|
|
)
|
|
x[..., split:] = model_sampling.noise_scaling(
|
|
audio_sigmas[0],
|
|
noise[..., split:],
|
|
latent_image[..., split:],
|
|
self.max_denoise(model_wrap, audio_sigmas),
|
|
)
|
|
s_in = x.new_ones([x.shape[0]])
|
|
held_v, held_a = x[..., :split], x[..., split:]
|
|
|
|
# a frozen stream is never stepped, so it needs no noise source
|
|
opts_v = (
|
|
None
|
|
if frozen_v
|
|
else self.stream_options(x, split, True, seed, model_sampling)
|
|
)
|
|
opts_a = (
|
|
None
|
|
if frozen_a
|
|
else self.stream_options(x, split, False, seed, model_sampling)
|
|
)
|
|
|
|
gen_v = (
|
|
None if frozen_v else ancestral_stream(held_v, video_sigmas, opts_v)
|
|
)
|
|
gen_a = (
|
|
None if frozen_a else ancestral_stream(held_a, audio_sigmas, opts_a)
|
|
)
|
|
|
|
progress = trange(total_steps, disable=disable_pbar)
|
|
req_v = next(gen_v) if gen_v is not None else None
|
|
req_a = next(gen_a) if gen_a is not None else None
|
|
done_v = held_v if frozen_v else None
|
|
done_a = held_a if frozen_a else None
|
|
while True:
|
|
step = (req_v or req_a)[0]
|
|
sigma_v, xv = (
|
|
(req_v[1], req_v[2]) if req_v is not None else (0.0, held_v)
|
|
)
|
|
sigma_a, xa = (
|
|
(req_a[1], req_a[2]) if req_a is not None else (0.0, held_a)
|
|
)
|
|
sigma_v, sigma_a = float(sigma_v), float(sigma_a)
|
|
x = torch.cat([xv, xa], dim=-1)
|
|
|
|
x_in = x
|
|
if denoise_mask is not None:
|
|
x_in = self.apply_mask(
|
|
inner_model,
|
|
x,
|
|
noise,
|
|
latent_image,
|
|
denoise_mask,
|
|
sigma_v,
|
|
sigma_a,
|
|
split,
|
|
model_options,
|
|
)
|
|
|
|
# the model derives the audio timestep from the video sigma and the two flow
|
|
# shifts; pinning the video shift to 1.0 makes that mapping solvable for the
|
|
# audio shift that reproduces this evaluation's audio sigma exactly
|
|
transformer_options["minimax_h3_sigma_shift_video"] = 1.0
|
|
transformer_options["minimax_h3_sigma_shift_audio"] = (
|
|
audio_shift_for(sigma_v, sigma_a)
|
|
)
|
|
denoised = model_wrap(
|
|
x_in, s_in * sigma_v, model_options=model_options, seed=seed
|
|
)
|
|
|
|
# calculate_denoised applied the video sigma to the whole pack; recover the
|
|
# velocity and re-solve the audio stream's x0 at its own sigma
|
|
velocity = (x_in - denoised) / sigma_v
|
|
denoised = torch.cat(
|
|
[
|
|
denoised[..., :split],
|
|
x_in[..., split:] - velocity[..., split:] * sigma_a,
|
|
],
|
|
dim=-1,
|
|
)
|
|
|
|
if denoise_mask is not None:
|
|
denoised = denoised * denoise_mask + latent_image * (
|
|
1.0 - denoise_mask
|
|
)
|
|
|
|
if step is not None:
|
|
if callback is not None:
|
|
# Sampler.sample gets ComfyUI's callback directly, not k-diffusion's dict form
|
|
callback(step, denoised, x, total_steps)
|
|
progress.update(1)
|
|
|
|
if gen_v is not None:
|
|
try:
|
|
req_v = gen_v.send(denoised[..., :split])
|
|
except StopIteration as stop:
|
|
req_v, done_v = None, stop.value
|
|
if gen_a is not None:
|
|
try:
|
|
req_a = gen_a.send(denoised[..., split:])
|
|
except StopIteration as stop:
|
|
req_a, done_a = None, stop.value
|
|
if req_v is None and req_a is None:
|
|
break
|
|
if (
|
|
gen_v is not None
|
|
and gen_a is not None
|
|
and (req_v is None) != (req_a is None)
|
|
):
|
|
raise RuntimeError(
|
|
"dual sampler streams desynchronised; both schedules must have the same length"
|
|
)
|
|
progress.close()
|
|
|
|
return torch.cat(
|
|
[
|
|
model_sampling.inverse_noise_scaling(video_sigmas[-1], done_v),
|
|
model_sampling.inverse_noise_scaling(audio_sigmas[-1], done_a),
|
|
],
|
|
dim=-1,
|
|
)
|
|
|
|
def apply_mask(
|
|
self,
|
|
inner_model,
|
|
x,
|
|
noise,
|
|
latent_image,
|
|
denoise_mask,
|
|
sigma_v,
|
|
sigma_a,
|
|
split,
|
|
model_options,
|
|
):
|
|
"""KSamplerX0Inpaint's masked blend, with each stream pinned at its own sigma."""
|
|
if "denoise_mask_function" in model_options:
|
|
denoise_mask = model_options["denoise_mask_function"](
|
|
sigma_v,
|
|
denoise_mask,
|
|
extra_options={
|
|
"model": inner_model,
|
|
"sigmas": self.audio_sigmas,
|
|
},
|
|
)
|
|
latent_mask = 1.0 - denoise_mask
|
|
pinned = torch.cat(
|
|
[
|
|
inner_model.scale_latent_inpaint(
|
|
x=x[..., :split],
|
|
sigma=x.new_tensor([sigma_v]),
|
|
noise=noise[..., :split],
|
|
latent_image=latent_image[..., :split],
|
|
),
|
|
inner_model.scale_latent_inpaint(
|
|
x=x[..., split:],
|
|
sigma=x.new_tensor([sigma_a]),
|
|
noise=noise[..., split:],
|
|
latent_image=latent_image[..., split:],
|
|
),
|
|
],
|
|
dim=-1,
|
|
)
|
|
return x * denoise_mask + pinned * latent_mask
|