Compare commits

..

No commits in common. "master" and "v0.32.0" have entirely different histories.

48 changed files with 217 additions and 3755 deletions

View File

@ -37,7 +37,7 @@
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
- ComfyUI natively supports the latest open-source state of the art models.
- [Partner nodes](https://docs.comfy.org/tutorials/partner-nodes/overview#partner-nodes) provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
- It is available on Windows, Linux, and macOS, locally with our [desktop application](https://www.comfy.org/download), our [portable install](#installing) or on our [cloud](https://www.comfy.org/cloud).
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
- It integrates seamlessly into production pipelines with our API endpoints.
@ -74,7 +74,7 @@ See what ComfyUI can do with the [newer template workflows](https://comfy.org/wo
- [Image editing](https://comfy.org/workflows/tag/image-edit/): Flux Kontext, Flux.2 Klein, Qwen Image Edit, HiDream E1.1 and O1, OmniGen2, Boogu, JoyImage Edit, MageFlow Edit, and LongCat Image Edit.
- [Video generation](https://comfy.org/workflows/tag/video-generation/): Wan 2.1 and 2.2, LTX-Video 2 and 2.3, HunyuanVideo 1.5, Kandinsky 5 Video, CogVideoX, Cosmos Predict2, Bernini-R, SCAIL 2, and Mochi.
- [Audio and video generation](https://comfy.org/workflows/): MiniMax H3 and LTX-AV.
- [Audio generation](https://comfy.org/workflows/tag/text-to-audio/): ACE-Step 1.5, Stable Audio 3 and MiniMax Music 3
- [Audio generation](https://comfy.org/workflows/tag/text-to-audio/): ACE-Step 1.5 and Stable Audio 3.
- [3D and vision](https://comfy.org/workflows/): Hunyuan3D 2.1, TripoSplat, SeedVR2, SUPIR, Depth Anything 3, MoGe, SAM 3 and 3.1, RT-DETRv4, and BiRefNet.
- [Text generation](https://comfy.org/workflows/tag/text-generation/): Gemma 3 and 4, Qwen3, Qwen3.5, and Qwen3-VL, including multimodal inputs.
- Load complete checkpoints or separate diffusion models, VAEs, text encoders, LoRAs, ControlNets, adapters, and upscalers from supported model formats.

View File

@ -180,7 +180,6 @@ parser.add_argument("--disable-async-offload", action="store_true", help="Disabl
parser.add_argument("--disable-dynamic-vram", action="store_true", help="Disable dynamic VRAM and use estimate based model loading.")
parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable dynamic VRAM on systems where it's not enabled by default.")
parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.")
parser.add_argument("--disable-cuda-graphs", action="store_true", help="Disable CUDA graphs.")
parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")

View File

@ -957,11 +957,6 @@ class ACEAudio15(LatentFormat):
latent_dimensions = 1
temporal_downscale_ratio = 1764
class MiniMaxMusic3(LatentFormat):
latent_channels = 128
latent_dimensions = 1
temporal_downscale_ratio = 512
class ChromaRadiance(LatentFormat):
latent_channels = 3
spacial_downscale_ratio = 1

View File

@ -22,7 +22,6 @@ import torch
import torch.nn.functional as F
from einops import rearrange
from torch import nn
import comfy.model_management
from comfy.ldm.lightricks.model import get_timestep_embedding
from .causal_video_autoencoder import Encoder, processor
@ -75,12 +74,8 @@ def default_rope_dim_split(head_dim):
def rope_inv_freqs(dim, base=10000.0, device=None):
out_device = device
if not comfy.model_management.supports_fp64(device):
device = torch.device("cpu")
exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(dtype=torch.float32, device=out_device)
return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(torch.float32)
def _rope_tables(lengths, inv_freqs, device):

View File

@ -91,18 +91,6 @@ def _video_t_grid(n, origin):
return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
def _ref_t_span(blk):
# time-axis span a reference block occupies ahead of the target streams
kind = blk["kind"]
if kind == "image":
return 1.0
if kind == "audio":
return float(blk["ref_audio_t"])
if kind in ("video", "video_audio"):
return max(float(blk["ref_audio_t"]), sum(_video_t_spans(blk["latent_t"])))
return 0.0
def _audio_grid(cursor, t, w_low, w_high):
# channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0
g = torch.zeros(t * 2, 3, dtype=torch.float64)
@ -300,7 +288,7 @@ class FinalLayer(nn.Module):
class PackedLayout:
"""Static packed-sequence structure for one shape/conditioning signature."""
def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None):
def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None):
frame, w_grid = _frame_grid(latent_h, latent_w)
frame_rows = frame.shape[0]
@ -311,37 +299,29 @@ class PackedLayout:
img_pos, img_update = [], []
audio_pos, audio_update = [], []
cursor = text_len
row = text_len
target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
# refs pack between text and the targets, so the target timeline starts after their spans
cursor = float(text_len)
for blk in refs or ():
cursor += _ref_t_span(blk)
if keyframes:
# fl2va: keyframe cond rows right after text, sharing the target spatial grid;
# anchors count from the target timeline origin, FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
# fl2va: keyframe cond rows right after text, sharing the target spatial grid
for kf in keyframes:
cond_t = cursor + FRAME_RESCALE * kf["resolved_frame_index"]
video_latent = kf.get("latent")
if video_latent is not None:
vt = video_latent.shape[2]
n = vt * frame_rows
segments.append(("cond", n))
pos.append(_video_grid(vt, frame, cond_t))
img_pos.append(torch.arange(row, row + n))
img_update.append(torch.zeros(n, dtype=torch.bool))
row += n
audio_latent = kf.get("audio_latent")
if audio_latent is not None:
rt = audio_latent.shape[-1]
segments.append(("cond_audio", rt * 2))
pos.append(_audio_grid(cond_t, rt, *target_audio_w))
audio_pos.append(torch.arange(row, row + rt * 2))
audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
row += rt * 2
pixel_index = kf["resolved_frame_index"]
if pixel_index == 0:
cond_t = float(text_len)
elif frame_count is not None and pixel_index == frame_count - 1:
cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
else:
raise ValueError("only first/last keyframe anchors are supported")
g = torch.empty(frame_rows, 3, dtype=torch.float64)
g[:, 0] = cond_t
g[:, 1:] = frame
segments.append(("cond", frame_rows))
pos.append(g)
img_pos.append(torch.arange(row, row + frame_rows))
img_update.append(torch.zeros(frame_rows, dtype=torch.bool))
row += frame_rows
target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
if refs:
cursor = float(text_len)
for blk in refs:
@ -409,7 +389,7 @@ class PackedLayout:
self.audio_update = torch.cat(audio_update)
self.signature = (text_len, latent_t, latent_h, latent_w, audio_t)
# contiguous segment table (start, stop, kind)
# kinds: text / cond / cond_audio / ref_img / ref_audio / audio / video
# kinds: text / cond / ref_img / ref_audio / audio / video
# the packed sequence is uniform per segment in (modality tag, timestep class),
# except the text span (tag runs resolved at forward time from the presentation tags)
seg_abs = []
@ -549,7 +529,8 @@ class MiniMaxH3Model(nn.Module):
if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t):
layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t,
keyframes=payload.get("keyframes"),
refs=payload.get("refs"))
refs=payload.get("refs"),
frame_count=payload.get("frame_count"))
# model_base passes model_sampling.timestep(sigma) = sigma * 1000
shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
@ -562,14 +543,14 @@ class MiniMaxH3Model(nn.Module):
vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments)
has_aud_cond = any(k in ("cond_audio", "ref_audio") for _, _, k in layout.segments)
has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments)
seg_t = {"text": t_v, "video": t_v, "audio": t_a,
"cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
"cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)}
"ref_audio": max(t_a, aud_aug)}
unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set())
| ({seg_t["ref_audio"]} if has_aud_cond else set()))
t_row = {t: i for i, t in enumerate(unique_t)}
seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2}
seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
text_tags = payload.get("text_token_tags")
mod_segments = []

View File

@ -1,343 +0,0 @@
import dataclasses
import hashlib
import torch
from torch import nn
import comfy.model_management
import comfy.model_prefetch
import comfy.ops
import comfy.utils
from comfy.ldm.modules.attention import optimized_attention_for_device
from comfy.text_encoders.llama import Llama2_, Qwen3_8BConfig
from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS
CFG_SCALE = 1.5
CFG_TOP_K = 50
C0_VOCAB_SIZE = 16384
MAX_PROMPT_TOKENS = 5000
MAX_AUDIO_FRAMES = 9000
AUDIO_FRAMES_PER_SECOND = 25
def derive_seed(seed, *parts):
digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm")
digest.update(int(seed).to_bytes(8, "little", signed=False))
for part in parts:
value = str(part).encode("utf-8")
digest.update(len(value).to_bytes(4, "little"))
digest.update(value)
return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1)
def sample_topk(logits, top_k, generator):
values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9)
top_k = min(top_k, values.shape[-1])
threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None]
values = values.masked_fill(values < threshold, -float("inf"))
probabilities = torch.nan_to_num(torch.softmax(values, dim=-1), nan=0.0)
probabilities = probabilities / probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12)
return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1)
class RVQAttention(nn.Module):
def __init__(self, hidden_size, num_heads, merged_qkv, dtype, device, operations):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.merged_qkv = merged_qkv
if merged_qkv:
self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
else:
self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
def forward(self, x):
batch, length, hidden_size = x.shape
if self.merged_qkv:
q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
else:
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
q = q.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
k = k.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
v = v.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
mask = torch.full((length, length), torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype).triu_(1)
attention = optimized_attention_for_device(q.device, mask=True, small_input=True)
out = attention(q, k, v, self.num_heads, mask=mask, skip_reshape=True)
return self.o_proj(out)
class RVQRMSNorm(nn.Module):
def __init__(self, hidden_size, dtype, device):
super().__init__()
self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device))
def forward(self, x):
return torch.nn.functional.rms_norm(x, (x.shape[-1],), comfy.ops.cast_to_input(self.weight, x), 1e-6)
class RVQMLP(nn.Module):
def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations):
super().__init__()
self.merged_mlp = merged_mlp
if merged_mlp:
self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
else:
self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device)
def forward(self, x):
if self.merged_mlp:
return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu")
return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
class RVQDecoderBlock(nn.Module):
def __init__(self, hidden_size, num_heads, intermediate_size, merged_qkv, merged_mlp, dtype, device, operations):
super().__init__()
self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device)
self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations)
self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device)
self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, dtype, device, operations)
def forward(self, x):
x = x + self.self_attn(self.input_layernorm(x))
return x + self.mlp(self.post_attention_layernorm(x))
class RVQDepthDecoder(nn.Module):
def __init__(self, config, dtype, device, operations):
super().__init__()
hidden_size = int(config["hidden_size"])
audio_vocab_size = int(config["audio_vocab_size"])
merged_qkv = config.get("decoder_merged_qkv", False)
merged_mlp = config.get("decoder_merged_mlp", False)
num_codebooks = int(config["audio_num_codebooks"])
self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device)
self.audio_heads = nn.ModuleList([
operations.Linear(hidden_size, audio_vocab_size, bias=False, dtype=dtype, device=device)
for _ in range(num_codebooks - 1)
])
self.layers = nn.ModuleList([
RVQDecoderBlock(
hidden_size,
int(config["decoder_num_heads"]),
int(config["decoder_intermediate_size"]),
merged_qkv,
merged_mlp,
dtype,
device,
operations,
)
for _ in range(int(config["decoder_num_layers"]))
])
self.norm = RVQRMSNorm(hidden_size, dtype, device)
def forward(self, sequence):
positions = torch.arange(sequence.shape[1], device=sequence.device)
x = sequence + self.pos_embedding(positions, out_dtype=sequence.dtype).unsqueeze(0)
for layer in self.layers:
x = layer(x)
return self.norm(x)
class MiniMaxMusic3AR(nn.Module):
def __init__(self, config, dtype, device, operations):
super().__init__()
config_fields = {field.name for field in dataclasses.fields(Qwen3_8BConfig)}
qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields})
qwen_config.lm_head = False
qwen_config.fixed_kv = True
self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations)
self.model.prefetch_dynamic_vbars = True
self.model.graph_dynamic_vbar_blocks = True
self.model.lm_head = operations.Linear(qwen_config.hidden_size, qwen_config.vocab_size, bias=False, dtype=dtype, device=device)
self.model.lm_head_pruned = operations.Linear(qwen_config.hidden_size, C0_VOCAB_SIZE + 1, bias=False, dtype=dtype, device=device)
self.model.embed_tokens_prefill = operations.Embedding(AUDIO_CODE_OFFSET, qwen_config.hidden_size, dtype=dtype, device=device)
self.model.embed_tokens_audio = operations.Embedding(C0_VOCAB_SIZE, qwen_config.hidden_size, dtype=dtype, device=device)
self.model.pruned_lm_head = None
self.model.pruned_embedding = None
self.model.audio_extra_embedding = operations.Embedding(
int(config["audio_vocab_size"]) * (int(config["audio_num_codebooks"]) - 1),
qwen_config.hidden_size,
dtype=dtype,
device=device,
)
self.model.audio_decoder = RVQDepthDecoder(config, dtype, device, operations)
self.audio_vocab_size = int(config["audio_vocab_size"])
self.num_codebooks = int(config["audio_num_codebooks"])
self.embedding_scale = self.num_codebooks ** -0.5
def _guided_c0(self, logits, cfg_scale, top_k):
conditioned = logits[0:1].float()
unconditioned = logits[1:2].float()
guided = unconditioned + (conditioned - unconditioned) * cfg_scale
threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None]
return guided.masked_fill(conditioned < threshold, -float("inf"))
def _depth_codes(self, hidden, c0, c0_embed, generator, execution_dtype, cfg_scale, top_k):
decoder = self.model.audio_decoder
sequence = [decoder.projection(hidden).unsqueeze(1)]
sequence.append(decoder.projection(c0_embed).unsqueeze(1))
codes = [c0]
hidden_parts = []
for index in range(1, self.num_codebooks):
out = decoder(torch.cat(sequence, dim=1))[:, -1]
hidden_parts.append(out[:1].detach())
logits = decoder.audio_heads[index - 1](out)
conditioned = logits[:1].float()
unconditioned = logits[1:2].float()
code = sample_topk(unconditioned + (conditioned - unconditioned) * cfg_scale, top_k, generator).repeat(2)
codes.append(code)
if index < self.num_codebooks - 1:
embedding = self.model.audio_extra_embedding(
code + (index - 1) * self.audio_vocab_size,
out_dtype=execution_dtype,
)
sequence.append(decoder.projection(embedding).unsqueeze(1))
return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1)
def _embed_c0(self, codes, execution_dtype):
if self.model.pruned_embedding:
return self.model.embed_tokens_audio(codes, out_dtype=execution_dtype)
return self.model.embed_tokens(codes + AUDIO_CODE_OFFSET, out_dtype=execution_dtype)
def _embed_audio_frame(self, codes, execution_dtype):
c0 = self._embed_c0(codes[:, 0], execution_dtype)
offsets = torch.arange(self.num_codebooks - 1, device=codes.device) * self.audio_vocab_size
extra = self.model.audio_extra_embedding(codes[:, 1:] + offsets.unsqueeze(0), out_dtype=execution_dtype).sum(dim=1)
return ((c0 + extra) * self.embedding_scale).unsqueeze(1)
def _sample_c0(self, hidden, cfg_scale, top_k, generator, vocab_mask):
if self.model.pruned_lm_head:
guided = self._guided_c0(self.model.lm_head_pruned(hidden).float(), cfg_scale, top_k)
code = sample_topk(guided, top_k, generator)
stop_token = 0
offset = 1
else:
logits = self.model.lm_head(hidden).float()
stop_token = SPECIAL_TOKEN_IDS["<|audio_end|>"]
logits = logits.masked_fill(vocab_mask, -float("inf"))
guided = self._guided_c0(logits, cfg_scale, top_k).masked_fill(vocab_mask, -float("inf"))
code = sample_topk(guided, top_k, generator)
offset = AUDIO_CODE_OFFSET
return torch.where(code == stop_token, 0, code - offset), code, stop_token
def generate(self, input_ids, seed, max_audio_frames, device, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K):
prompt_tokens = int(input_ids.shape[1])
if prompt_tokens > MAX_PROMPT_TOKENS:
raise ValueError(f"MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is {MAX_PROMPT_TOKENS}")
input_ids = input_ids.to(device)
if comfy.model_management.should_use_bf16(device):
execution_dtype = torch.bfloat16
else:
execution_dtype = torch.float32
unconditioned = input_ids.clone()
unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"]
text_ids = torch.cat((input_ids, unconditioned), dim=0)
if self.model.pruned_embedding:
text_embeds = self.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype)
else:
text_embeds = self.model.embed_tokens(text_ids, out_dtype=execution_dtype)
decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES)
past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype)
output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype)
last_hidden = output[0][:, -1]
past = output[2]
generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar"))
decoder = self.model.audio_decoder
depth_io = {
"hidden": torch.empty_like(last_hidden),
"c0": torch.empty((last_hidden.shape[0],), dtype=torch.long, device=device),
"c0_embed": torch.empty_like(last_hidden),
"codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device),
"depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device),
}
decoder._comfy_cross_step_state = depth_io
comfy.model_management._register_cross_step(decoder)
hidden_frames = []
pending_code = None
stop_token = None
pending_event = None
pending_hidden = None
progress = comfy.utils.ProgressBar(decode_limit)
cuda_device = torch.device(device).type == "cuda"
vocab_mask = None
if not self.model.pruned_lm_head:
vocab_mask = torch.ones(self.model.vocab_size, dtype=torch.bool, device=device)
vocab_mask[AUDIO_CODE_OFFSET:AUDIO_CODE_OFFSET + C0_VOCAB_SIZE] = False
vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False
for frame_index in comfy.utils.model_trange(decode_limit + 1, desc="AR sampling"):
comfy.model_management.throw_exception_if_processing_interrupted()
if pending_code is not None:
if pending_event is not None:
pending_event.synchronize()
if int(pending_code.item()) == stop_token:
pending_hidden = None
break
if pending_hidden is not None:
hidden_frames.append(pending_hidden)
progress.update_absolute(len(hidden_frames))
if len(hidden_frames) >= decode_limit:
break
c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
if pending_code is None:
pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device)
if cuda_device:
pending_event = torch.cuda.Event()
pending_code.copy_(code_or_stop, non_blocking=cuda_device)
if pending_event is not None:
pending_event.record()
c0 = c0.repeat(2)
c0_embed = self._embed_c0(c0, execution_dtype)
depth_io["hidden"].copy_(last_hidden)
depth_io["c0"].copy_(c0)
depth_io["c0_embed"].copy_(c0_embed)
def depth_core():
codes, depth_hidden = self._depth_codes(
depth_io["hidden"], depth_io["c0"], depth_io["c0_embed"], generator, execution_dtype, cfg_scale, top_k
)
depth_io["codes"].copy_(codes)
depth_io["depth_hidden"].copy_(depth_hidden)
depth_queue = comfy.model_prefetch.make_prefetch_queue(
[[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True}
)
comfy.model_prefetch.prefetch_queue_pop(
depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator
)
comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None)
feedback_codes = depth_io["codes"]
depth_hidden = depth_io["depth_hidden"]
frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1)
if frame_index > 0:
pending_hidden = frame_hidden[0].clone()
feedback = self._embed_audio_frame(feedback_codes, execution_dtype)
output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
last_hidden = output[0][:, -1]
past = output[2]
if pending_hidden is not None and len(hidden_frames) < decode_limit:
if pending_event is not None:
pending_event.synchronize()
if int(pending_code.item()) != stop_token:
hidden_frames.append(pending_hidden)
if not hidden_frames:
raise ValueError("MiniMax Music3 generated zero audio frames")
return torch.stack(hidden_frames).to(device="cpu")

View File

@ -1,137 +0,0 @@
import math
import torch
from torch import nn
import comfy.ops
def snake(x, alpha):
shape = x.shape
flat = x.reshape(shape[0], shape[1], -1)
alpha = comfy.ops.cast_to_input(alpha, flat)
flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2)
return flat.reshape(shape)
class Snake1d(nn.Module):
def __init__(self, channels, dtype, device):
super().__init__()
self.alpha = nn.Parameter(torch.empty(1, channels, 1, dtype=dtype, device=device))
def forward(self, x):
return snake(x, self.alpha)
def _weight_norm_conv(operations, *args, **kwargs):
return nn.utils.parametrizations.weight_norm(operations.Conv1d(*args, **kwargs))
def _weight_norm_conv_transpose(operations, *args, **kwargs):
return nn.utils.parametrizations.weight_norm(operations.ConvTranspose1d(*args, **kwargs))
class ResidualUnit(nn.Module):
def __init__(self, dim, dilation, dtype, device, operations):
super().__init__()
padding = 3 * dilation
self.block = nn.Sequential(
Snake1d(dim, dtype, device),
_weight_norm_conv(
operations,
dim,
dim,
kernel_size=7,
dilation=dilation,
padding=padding,
dtype=dtype,
device=device,
),
Snake1d(dim, dtype, device),
_weight_norm_conv(operations, dim, dim, kernel_size=1, dtype=dtype, device=device),
)
def forward(self, x):
residual = self.block(x)
if residual.shape[-1] != x.shape[-1]:
padding = (x.shape[-1] - residual.shape[-1]) // 2
x = x[..., padding:x.shape[-1] - padding]
return x + residual
class DecoderBlock(nn.Module):
def __init__(self, input_dim, output_dim, stride, dtype, device, operations):
super().__init__()
self.block = nn.Sequential(
Snake1d(input_dim, dtype, device),
_weight_norm_conv_transpose(
operations,
input_dim,
output_dim,
kernel_size=2 * stride,
stride=stride,
padding=math.ceil(stride / 2),
dtype=dtype,
device=device,
),
ResidualUnit(output_dim, 1, dtype, device, operations),
ResidualUnit(output_dim, 3, dtype, device, operations),
ResidualUnit(output_dim, 9, dtype, device, operations),
)
def forward(self, x):
return self.block(x)
class Decoder(nn.Module):
def __init__(self, dtype, device, operations):
super().__init__()
layers = [
_weight_norm_conv(
operations,
1024,
1536,
kernel_size=7,
padding=3,
dtype=dtype,
device=device,
)
]
channels = 1536
output_dim = channels
for index, stride in enumerate((8, 8, 4, 2)):
input_dim = channels // (2 ** index)
output_dim = channels // (2 ** (index + 1))
layers.append(DecoderBlock(input_dim, output_dim, stride, dtype, device, operations))
layers.extend((
Snake1d(output_dim, dtype, device),
_weight_norm_conv(
operations,
output_dim,
1,
kernel_size=7,
padding=3,
dtype=dtype,
device=device,
),
nn.Tanh(),
))
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)
class MiniMaxMusic3DAV(nn.Module):
def __init__(self, dtype=None, device=None, operations=None):
super().__init__()
self.dec_in_proj = operations.Conv1d(64, 1024, kernel_size=1, dtype=dtype, device=device)
self.decoder = Decoder(dtype, device, operations)
def decode(self, latent):
batch, _, frames = latent.shape
folded = latent.reshape(batch * 2, 64, frames)
waveform = self.decoder(self.dec_in_proj(folded))
return waveform.reshape(batch, 2, -1)
forward = decode

View File

@ -1,213 +0,0 @@
import math
import torch
from torch import nn
import comfy.model_management
import comfy.ops
import comfy.quant_ops
from comfy.ldm.modules.attention import optimized_attention_for_device
MAX_CONDITION_FRAMES = 200
CONDITION_HOP_FRAMES = 100
def latent_length(audio_frames):
return max(1, int(audio_frames * 44100 / 24000 * 960 / 512))
class FourierFeatures(nn.Module):
def __init__(self, in_features, out_features, dtype, device):
super().__init__()
self.weight = nn.Parameter(torch.empty(out_features // 2, in_features, dtype=dtype, device=device))
def forward(self, value):
weight = comfy.ops.cast_to_input(self.weight, value)
features = 2.0 * math.pi * value @ weight.T
return torch.cat((features.cos(), features.sin()), dim=-1)
class LayerNorm(nn.Module):
def __init__(self, dim, dtype, device):
super().__init__()
self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
self.register_buffer("beta", torch.empty(dim, dtype=dtype, device=device))
def forward(self, x):
return torch.nn.functional.layer_norm(
x,
(x.shape[-1],),
comfy.ops.cast_to_input(self.gamma, x),
comfy.ops.cast_to_input(self.beta, x),
)
class RotaryEmbedding(nn.Module):
def __init__(self, dim, dtype, device):
super().__init__()
self.register_buffer("inv_freq", torch.empty(dim // 2, dtype=dtype, device=device))
def forward_from_seq_len(self, length, device, dtype):
positions = torch.arange(length, device=device, dtype=torch.float32)
frequencies = torch.outer(positions, comfy.ops.cast_to_input(self.inv_freq, positions))
frequencies = frequencies.to(dtype)
cos, sin = frequencies.cos(), frequencies.sin()
return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(1, 1, length, frequencies.shape[-1], 2, 2)
def _apply_rope(x, rotation_matrix):
x_dtype = x.dtype
x = x.reshape(*x.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).to(rotation_matrix.dtype)
x = rotation_matrix[..., 0] * x[..., 0] + rotation_matrix[..., 1] * x[..., 1]
return x.movedim(-1, -2).flatten(-2).to(x_dtype)
class Attention(nn.Module):
def __init__(self, dim, dim_heads, dtype, device, operations):
super().__init__()
self.num_heads = dim // dim_heads
self.dim_heads = dim_heads
self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device)
self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
def forward(self, x, rotation_matrix):
batch, length, dim = x.shape
q, k, v = self.to_qkv(x).chunk(3, dim=-1)
q = q.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
k = k.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
v = v.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
rotary_dims = rotation_matrix.shape[-3] * 2
if comfy.model_management.in_training:
q = torch.cat((_apply_rope(q[..., :rotary_dims], rotation_matrix), q[..., rotary_dims:]), dim=-1)
k = torch.cat((_apply_rope(k[..., :rotary_dims], rotation_matrix), k[..., rotary_dims:]), dim=-1)
else:
rotated_q, rotated_k = comfy.quant_ops.ck.apply_rope_split_half(q[..., :rotary_dims], k[..., :rotary_dims], rotation_matrix)
q = torch.cat((rotated_q, q[..., rotary_dims:]), dim=-1)
k = torch.cat((rotated_k, k[..., rotary_dims:]), dim=-1)
attention = optimized_attention_for_device(q.device)
out = attention(q, k, v, self.num_heads, skip_reshape=True)
return self.to_out(out)
class GLU(nn.Module):
def __init__(self, dim, inner_dim, dtype, device, operations):
super().__init__()
self.proj = operations.Linear(dim, inner_dim * 2, dtype=dtype, device=device)
def forward(self, x):
value, gate = self.proj(x).chunk(2, dim=-1)
return value * torch.nn.functional.silu(gate)
class FeedForward(nn.Module):
def __init__(self, dim, inner_dim, dtype, device, operations):
super().__init__()
self.ff = nn.Sequential(
GLU(dim, inner_dim, dtype, device, operations),
nn.Identity(),
operations.Linear(inner_dim, dim, dtype=dtype, device=device),
)
def forward(self, x):
return self.ff(x)
class TransformerBlock(nn.Module):
def __init__(self, dim, dim_heads, inner_dim, dtype, device, operations):
super().__init__()
self.pre_norm = LayerNorm(dim, dtype, device)
self.self_attn = Attention(dim, dim_heads, dtype, device, operations)
self.ff_norm = LayerNorm(dim, dtype, device)
self.ff = FeedForward(dim, inner_dim, dtype, device, operations)
def forward(self, x, rotation_matrix):
x = x + self.self_attn(self.pre_norm(x), rotation_matrix)
return x + self.ff(self.ff_norm(x))
class ContinuousTransformer(nn.Module):
def __init__(self, dtype, device, operations):
super().__init__()
self.project_in = operations.Linear(2304, 2048, bias=False, dtype=dtype, device=device)
self.project_out = operations.Linear(2048, 128, bias=False, dtype=dtype, device=device)
self.rotary_pos_emb = RotaryEmbedding(32, dtype, device)
self.layers = nn.ModuleList([
TransformerBlock(2048, 64, 8192, dtype, device, operations)
for _ in range(36)
])
def forward(self, x, timestep_embedding):
x = self.project_in(x)
x = torch.cat((timestep_embedding.unsqueeze(1), x), dim=1)
rotation_matrix = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], x.device, x.dtype)
for layer in self.layers:
x = layer(x, rotation_matrix)
return self.project_out(x[:, 1:])
class DiffusionTransformer(nn.Module):
def __init__(self, dtype, device, operations):
super().__init__()
self.transformer = ContinuousTransformer(dtype, device, operations)
self.timestep_features = FourierFeatures(1, 256, dtype, device)
self.to_timestep_embed = nn.Sequential(
operations.Linear(256, 2048, dtype=dtype, device=device),
nn.SiLU(),
operations.Linear(2048, 2048, dtype=dtype, device=device),
)
self.preprocess_conv = operations.Conv1d(2304, 2304, 1, bias=False, dtype=dtype, device=device)
self.postprocess_conv = operations.Conv1d(128, 128, 1, bias=False, dtype=dtype, device=device)
def forward(self, x, timestep, condition):
full = torch.cat((x, torch.zeros_like(x), condition), dim=1)
full = self.preprocess_conv(full) + full
timestep_features = self.timestep_features(timestep[:, None]).to(dtype=x.dtype)
timestep_embedding = self.to_timestep_embed(timestep_features)
out = self.transformer(full.transpose(1, 2), timestep_embedding).transpose(1, 2)
return self.postprocess_conv(out) + out
class MiniMaxMusic3DiT(nn.Module):
def __init__(self, dtype=None, device=None, operations=None, **kwargs):
super().__init__()
self.dtype = dtype
self.latent_conditioners = nn.Sequential(
operations.Conv1d(4096, 2048, kernel_size=3, padding=1, dtype=dtype, device=device)
)
self.diffusion_transformer = DiffusionTransformer(dtype, device, operations)
self.cond_layer_logits = nn.Parameter(torch.empty(8, dtype=dtype, device=device))
self.cond_layer_scale = nn.Parameter(torch.empty(1, dtype=dtype, device=device))
def aligned_condition(self, hidden):
frames = hidden.shape[1]
hidden = hidden.transpose(1, 2).reshape(hidden.shape[0], 8, 4096, frames)
weights = torch.softmax(comfy.ops.cast_to_input(self.cond_layer_logits, hidden), dim=0)
hidden = torch.einsum("blht,l->bht", hidden, weights)
hidden = comfy.ops.cast_to_input(self.cond_layer_scale, hidden) * hidden
condition = self.latent_conditioners(hidden)
return torch.nn.functional.interpolate(condition, size=latent_length(frames), mode="nearest")
def forward(self, x, timestep, context, conditioning_scale, **kwargs):
condition = self.aligned_condition(context)
condition = condition * conditioning_scale[:, :1, :1]
if condition.shape[-1] < x.shape[-1]:
condition = torch.nn.functional.pad(condition, (0, x.shape[-1] - condition.shape[-1]))
else:
condition = condition[..., :x.shape[-1]]
window = latent_length(MAX_CONDITION_FRAMES)
if x.shape[-1] <= window:
return -self.diffusion_transformer(x, timestep, condition)
output = torch.zeros_like(x)
count = torch.zeros((1, 1, x.shape[-1]), device=x.device, dtype=x.dtype)
hop = latent_length(CONDITION_HOP_FRAMES)
start = 0
while start < x.shape[-1]:
end = min(start + window, x.shape[-1])
output[..., start:end] -= self.diffusion_transformer(x[..., start:end], timestep, condition[..., start:end])
count[..., start:end] += 1
if end == x.shape[-1]:
break
start += hop
return output / count

View File

@ -1,70 +0,0 @@
import re
SPECIAL_TOKEN_IDS = {
"<|im_start|>": 151644,
"<|im_end|>": 151645,
"<|audio_cfg|>": 151654,
"<|audio_start|>": 151669,
"<|audio_end|>": 151670,
"<|caption_start|>": 151671,
"<|caption_end|>": 151672,
"<|lyrics_start|>": 151673,
"<|lyrics_end|>": 151674,
}
AUDIO_CODE_OFFSET = 151675
_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>")
_LYRIC_TAG_RE = re.compile(r"\s*(\[[^\]]+\])\s*")
def _remove_markdown_format(text):
lines = []
for raw_line in text.splitlines():
line = re.sub(r"^\s{0,3}#{1,6}\s+", "", raw_line)
line = re.sub(r"^\s*[*+-]\s+", "", line)
while "**" in line:
updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line)
if updated == line:
break
line = updated
line = re.sub(r"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
lines.append(line.rstrip())
text = "\n".join(lines)
text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
return text.replace("", "").replace(" ", "")
def clean_caption(caption):
def replace_special(match):
inner = match.group(1).strip()
parts = inner.split(None, 1)
return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner
text = _SPECIAL_TAG_RE.sub(replace_special, caption)
text = _remove_markdown_format(text)
return re.sub(r"\n{2,}", "\n", text)
def normalize_lyrics(lyrics):
parts = _LYRIC_TAG_RE.split(lyrics)
text = "\n".join(part.lower() if part.startswith("[") else part for part in parts if part)
text = text.replace(" ^ ", "\n")
return f"[start]\n{text}"
def build_prompt(caption, lyrics):
return (
"<|im_start|><|caption_start|>"
f"{clean_caption(caption)}"
"<|caption_end|><|lyrics_start|>"
f"{normalize_lyrics(lyrics)}"
"<|lyrics_end|><|im_end|><|audio_start|>"
)
def validate_tokenizer(tokenizer):
for token, expected in SPECIAL_TOKEN_IDS.items():
token_id = tokenizer.convert_tokens_to_ids(token)
if token_id != expected:
raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}: expected {expected}, got {token_id}")

View File

@ -22,7 +22,6 @@ import torch
import logging
import comfy.ldm.lightricks.av_model
import comfy.ldm.minimax.model
import comfy.ldm.minimax_music.dit
import comfy.nested_tensor
import comfy.ldm.lightricks.symmetric_patchifier
import comfy.context_windows
@ -2165,13 +2164,13 @@ class MiniMaxH3(BaseModel):
keyframes = kwargs.get("minimax_keyframes", None)
if keyframes is not None:
payload["keyframes"] = keyframes
payload["cond_video_latents"] = [kf["latent"] for kf in keyframes if kf.get("latent") is not None]
payload["cond_audio_latents"] = [kf["audio_latent"] for kf in keyframes if kf.get("audio_latent") is not None]
payload["frame_count"] = kwargs.get("minimax_frame_count", None)
payload["cond_video_latents"] = [kf["latent"] for kf in keyframes]
refs = kwargs.get("minimax_refs", None)
if refs is not None:
payload["refs"] = refs
payload["cond_video_latents"] = payload.get("cond_video_latents", []) + [r["latent"] for r in refs if "latent" in r]
payload["cond_audio_latents"] = payload.get("cond_audio_latents", []) + [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
payload["cond_video_latents"] = [r["latent"] for r in refs if "latent" in r]
payload["cond_audio_latents"] = [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
if kwargs.get("minimax_visual_cond_noise_aug", None) is not None:
payload["visual_cond_noise_aug"] = kwargs["minimax_visual_cond_noise_aug"]
if kwargs.get("minimax_audio_cond_noise_aug", None) is not None:
@ -2185,7 +2184,7 @@ class MiniMaxH3(BaseModel):
payload["layout"] = comfy.ldm.minimax.model.PackedLayout(
cross_attn.shape[1], vs[2], (vs[3] + 1) // 2 * 2, (vs[4] + 1) // 2 * 2,
latent_shapes[1][-1], keyframes=payload.get("keyframes"),
refs=payload.get("refs"))
refs=payload.get("refs"), frame_count=payload.get("frame_count"))
out['minimax_payload'] = comfy.conds.CONDConstant(payload)
return out
@ -2338,18 +2337,6 @@ class ACEStep15(BaseModel):
out['refer_audio'] = comfy.conds.CONDRegular(refer_audio)
return out
class MiniMaxMusic3(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax_music.dit.MiniMaxMusic3DiT)
def process_timestep(self, timestep, **kwargs):
return 1.0 - timestep
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
out["conditioning_scale"] = comfy.conds.CONDRegular(kwargs["conditioning_scale"])
return out
class Omnigen2(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.omnigen.omnigen2.OmniGen2Transformer2DModel)

View File

@ -44,13 +44,6 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
def detect_unet_config(state_dict, key_prefix, metadata=None):
state_dict_keys = list(state_dict.keys())
if (
'{}cond_layer_logits'.format(key_prefix) in state_dict_keys
and '{}latent_conditioners.0.weight'.format(key_prefix) in state_dict_keys
and '{}diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix) in state_dict_keys
):
return {"audio_model": "minimax_music3"}
if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model
unet_config = {}
unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]
@ -837,10 +830,11 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["use_adaln_lora"] = True
dit_config["adaln_lora_dim"] = 256
dit_config["num_blocks"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
if dit_config["model_channels"] == 2048:
dit_config["num_blocks"] = 28
dit_config["num_heads"] = 16
elif dit_config["model_channels"] == 5120:
dit_config["num_blocks"] = 36
dit_config["num_heads"] = 40
if dit_config["in_channels"] == 16:

View File

@ -490,36 +490,28 @@ try:
except:
rocm_version = (6, -1)
def aotriton_supported():
"""Whether pytorch reports flash attention as usable on this gpu.
can_use_flash_attention() evaluates runtime eligibility for the given
parameters; on a ROCm build that includes checking the gpu arch against the
kernel images AOTriton was compiled for. Querying it avoids assuming where
those images live inside the torch install. The probe tensor is shaped and
typed to pass the unrelated SDPA checks, so False means no hardware support
rather than a rejected shape.
"""
try:
if not torch.backends.cuda.is_flash_attention_available(): # not built with flash attention
return False
q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
return torch.backends.cuda.can_use_flash_attention(params, False)
except (AttributeError, RuntimeError, TypeError) as e:
logging.warning("Could not query aotriton support: {}".format(e))
return False
def aotriton_supported(gpu_arch):
path = torch.__path__[0]
path = os.path.join(os.path.join(path, "lib"), "aotriton.images")
gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path))))
if gpu_arch in gfx:
return True
if "{}x".format(gpu_arch[:-1]) in gfx:
return True
if "{}xx".format(gpu_arch[:-2]) in gfx:
return True
return False
logging.info("AMD arch: {}".format(arch))
logging.info("ROCm version: {}".format(rocm_version))
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
if aotriton_supported(): # AMD efficient attention implementation depends on aotriton.
if aotriton_supported(arch): # AMD efficient attention implementation depends on aotriton.
if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much
if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]): # TODO: more arches, TODO: gfx950
ENABLE_PYTORCH_ATTENTION = True
if rocm_version >= (7, 0):
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
ENABLE_PYTORCH_ATTENTION = True
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
ENABLE_PYTORCH_ATTENTION = True
if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
SUPPORT_FP8_OPS = True
@ -1368,14 +1360,9 @@ STREAM_CAST_BUFFERS = {}
LARGEST_CASTED_WEIGHT = (None, 0)
STREAM_AIMDO_CAST_BUFFERS = {}
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
CROSS_STEP_STATE = weakref.WeakSet()
DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
# NOTE: devs/agents: this is temporary and will be removed in a future comfy. Not supported for custom node use.
def _register_cross_step(module):
CROSS_STEP_STATE.add(module)
def get_cast_buffer(offload_stream, device, size, ref):
global LARGEST_CASTED_WEIGHT
@ -1430,10 +1417,6 @@ def reset_cast_buffers():
mmap_obj.bounce()
DIRTY_MMAPS.clear()
for module in CROSS_STEP_STATE:
del module._comfy_cross_step_state
CROSS_STEP_STATE.clear()
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is not None and model.is_dynamic():

View File

@ -1887,29 +1887,8 @@ class ModelPatcherDynamic(ModelPatcher):
loading = self._load_list(for_dynamic=True, default_device=device_to)
loading.sort()
get_units = getattr(self.model, "get_dynamic_vram__units", None)
dynamic_units, last_dynamic_units = get_units() if get_units is not None else ([], [])
dynamic_units = list(dynamic_units)
last_dynamic_units = list(last_dynamic_units)
loading_by_module = {entry[-2]: entry for entry in loading}
loading = []
for unit in dynamic_units:
unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
for index, module in enumerate(modules):
loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
last_loading = []
for unit in last_dynamic_units:
unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
for index, module in enumerate(modules):
last_loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
loading.extend((*entry, None) for entry in loading_by_module.values())
loading.extend(last_loading)
v_block = None
for x in loading:
*_, module_mem, n, m, params, end_of_block = x
*_, module_mem, n, m, params = x
def set_dirty(item, dirty):
if dirty or not hasattr(item, "_v_signature"):
@ -2002,13 +1981,6 @@ class ModelPatcherDynamic(ModelPatcher):
move_weight_functions(m, device_to)
if hasattr(m, "_v"):
v_block = m._v if v_block is None else (v_block[0], v_block[1], max(v_block[2], m._v[1] + m._v[2] - v_block[1]))
if end_of_block is not None:
unit = end_of_block
(unit[0] if isinstance(unit, (list, tuple)) else unit)._v_block = v_block
v_block = None
for key, buf in self.model.named_buffers(recurse=True):
if key not in self.backup_buffers:
self.backup_buffers[key] = buf

View File

@ -1,19 +1,11 @@
import torch
import warnings
import weakref
import comfy_aimdo.model_vbar
from comfy.cli_args import args
import comfy.memory_management
import comfy.model_management
import comfy.ops
PREFETCH_QUEUES = []
GRAPH_MODULES = weakref.WeakSet()
GRAPH_WARMED_MODULES = weakref.WeakSet()
GRAPH_CAPTURE_STREAMS = {}
def cleanup_prefetched_modules(module, comfy_modules):
def cleanup_prefetched_modules(comfy_modules):
for s in comfy_modules:
prefetch = getattr(s, "_prefetch", None)
if prefetch is None:
@ -25,86 +17,39 @@ def cleanup_prefetched_modules(module, comfy_modules):
if prefetch["signature"] is not None:
comfy_aimdo.model_vbar.vbar_unpin(s._v)
delattr(s, "_prefetch")
if getattr(module, "_v_block_faulted", False):
comfy_aimdo.model_vbar.vbar_unpin(module._v_block)
del module._v_block_faulted
def _drop_graph(module):
graph = getattr(module, "_comfy_graph", None)
if graph is None:
return
# reset() through the bound method surfaces the allocator's benign
# "uncaptured free of a captured allocation" as catchable Python warnings;
# a plain del frees from the C++ dealloc path and spams stderr instead
with warnings.catch_warnings():
warnings.simplefilter("ignore")
graph["graph"].reset()
del module._comfy_graph
def cleanup_prefetch_queues():
global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS
global PREFETCH_QUEUES
for queue in PREFETCH_QUEUES:
for entry in queue:
if entry is None or not isinstance(entry, tuple):
continue
_, prefetch_state = entry
prefetched_module, comfy_modules = prefetch_state
comfy_modules = prefetch_state[1]
if comfy_modules is not None:
cleanup_prefetched_modules(prefetched_module, comfy_modules)
cleanup_prefetched_modules(comfy_modules)
PREFETCH_QUEUES = []
for module in GRAPH_MODULES:
_drop_graph(module)
GRAPH_MODULES.clear()
GRAPH_WARMED_MODULES.clear()
GRAPH_CAPTURE_STREAMS = {}
def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None):
enable_graph = enable_graph and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None
def prefetch_queue_pop(queue, device, module):
if queue is None:
if core is not None:
core()
return
capture_stream = None
if enable_graph:
capture_stream = GRAPH_CAPTURE_STREAMS.get(device)
if capture_stream is None:
capture_stream = torch.cuda.Stream(device=device)
GRAPH_CAPTURE_STREAMS[device] = capture_stream
signature = None
graph_hit = False
graph = getattr(module, "_comfy_graph", None) if enable_graph else None
if graph is not None:
signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
if signature is not None:
module._v_block_faulted = True
graph_hit = comfy_aimdo.model_vbar.vbar_signature_compare(signature, graph["signature"])
consumed = queue.pop(0)
if consumed is not None:
offload_stream, prefetch_state = consumed
if offload_stream is not None:
offload_stream.wait_stream(comfy.model_management.current_stream(device))
prefetched_module, comfy_modules = prefetch_state
_, comfy_modules = prefetch_state
if comfy_modules is not None:
cleanup_prefetched_modules(prefetched_module, comfy_modules)
cleanup_prefetched_modules(comfy_modules)
if graph_hit:
queue[0] = (None, (module, []))
graph["graph"].replay()
return
fully_faulted = False
prefetch = queue[0]
if prefetch is not None:
comfy_modules = []
prefetch_modules = prefetch if isinstance(prefetch, (list, tuple)) else (prefetch,)
for root in prefetch_modules:
for s in root.modules():
if hasattr(s, "_v"):
comfy_modules.append(s)
for s in prefetch.modules():
if hasattr(s, "_v"):
comfy_modules.append(s)
registerable_size = 0
for s in comfy_modules:
@ -114,42 +59,11 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
if lowvram_fn is not None:
registerable_size += lowvram_fn.memory_required()
offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True)
offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True)
if not comfy.model_management.args.fast_disk:
comfy.model_management.ensure_pin_registerable(registerable_size)
comfy.model_management.sync_stream(device, offload_stream)
if fully_faulted and dtype is not None:
for comfy_module in comfy_modules:
comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False)
queue[0] = (offload_stream, (module, comfy_modules))
if core is not None:
if enable_graph and fully_faulted and module in GRAPH_WARMED_MODULES:
if signature is None:
signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
if signature is not None:
module._v_block_faulted = True
if signature is not None:
_drop_graph(module)
graph = torch.cuda.CUDAGraph()
if generator is not None:
graph.register_generator_state(generator)
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"):
core()
comfy.model_management.current_stream(device).wait_stream(capture_stream)
graph.replay()
module._comfy_graph = {"graph": graph, "signature": signature}
GRAPH_MODULES.add(module)
return
if capture_stream is None:
core()
else:
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.stream(capture_stream):
core()
comfy.model_management.current_stream(device).wait_stream(capture_stream)
GRAPH_WARMED_MODULES.add(module)
queue[0] = (offload_stream, (prefetch, comfy_modules))
def make_prefetch_queue(queue, device, transformer_options):
if (not transformer_options.get("prefetch_dynamic_vbars", False)

View File

@ -123,12 +123,10 @@ def materialize_meta_param(s, param_keys):
# FIXME: add n=1 cache hit fast path
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False):
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking):
offload_stream = None
cast_buffer = None
cast_buffer_offset = 0
if return_faulted:
fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias"))
def ensure_offload_stream(module, required_size, check_largest):
nonlocal offload_stream
@ -165,8 +163,6 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
for s in comfy_modules:
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
if return_faulted and (signature is None or not resident):
fully_faulted = False
prefetch = {
"signature": signature,
"resident": resident,
@ -259,12 +255,10 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
prefetch["needs_cast"] = needs_cast
s._prefetch = prefetch
if return_faulted:
return offload_stream, fully_faulted
return offload_stream
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True):
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant):
prefetch = getattr(s, "_prefetch", None)
@ -304,7 +298,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
tensor = tensor.dequantize()
return tensor
if (return_weights and orig.dtype != dtype) or len(fns) > 0:
if orig.dtype != dtype or len(fns) > 0:
x = to_dequant(x, dtype)
if not resident and lowvram_fn is not None:
x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
@ -331,7 +325,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
if prefetch["signature"] is not None:
prefetch["resident"] = True
return (weight, bias) if return_weights else None
return weight, bias
def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
@ -1639,7 +1633,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
self.norm_type, self.scale_grad_by_freq, self.sparse)
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
x = x.to(dtype=target_dtype)
if scale is not None:
if scale is not None and scale != 1.0:
x = x * scale.to(dtype=target_dtype)
return x

View File

@ -37,11 +37,6 @@ def prepare_noise(latent_image, seed, noise_inds=None):
return noises
def prepare_empty_noise(latent_image):
if latent_image.is_nested:
return comfy.nested_tensor.NestedTensor([torch.zeros_like(t, device="cpu") for t in latent_image.unbind()])
return torch.zeros_like(latent_image, device="cpu")
def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None, downscale_ratio_temporal=None):
if latent_image.is_nested:
return latent_image

View File

@ -25,7 +25,6 @@ import comfy.ldm.cogvideo.vae
import comfy.ldm.hunyuan_video.vae
import comfy.ldm.mmaudio.vae.autoencoder
import comfy.ldm.audio.vae_sa3
import comfy.ldm.minimax_music.dav
import comfy.pixel_space_convert
import comfy.weight_adapter
import yaml
@ -33,7 +32,6 @@ import math
import os
import comfy.utils
import comfy.ops
from . import clip_vision
from . import gligen
@ -76,7 +74,6 @@ import comfy.text_encoders.longcat_image
import comfy.text_encoders.qwen35
import comfy.text_encoders.qwen3vl
import comfy.text_encoders.minimax
import comfy.text_encoders.minimax_music
import comfy.ldm.minimax.vae
import comfy.ldm.minimax.audio_vae
import comfy.text_encoders.boogu
@ -518,22 +515,7 @@ class VAE:
self.audio_sample_rate = 44100
if config is None:
if "dec_in_proj.weight" in sd and "decoder.model.0.weight_g" in sd: # MiniMax Music3 DAV
self.first_stage_model = comfy.ldm.minimax_music.dav.MiniMaxMusic3DAV(operations=comfy.ops.disable_weight_init)
self.latent_channels = 128
self.output_channels = 2
self.upscale_ratio = 512
self.downscale_ratio = 512
self.latent_dim = 1
self.process_output = lambda audio: audio
self.process_input = lambda audio: audio
self.working_dtypes = [torch.float32]
self.disable_offload = True
self.memory_used_decode = lambda shape, dtype: (shape[-1] * 512 * 1400 + 800_000_000) * model_management.dtype_size(dtype)
def _no_encode(*args, **kwargs):
raise RuntimeError("MiniMax Music3 DAV cannot encode audio")
self.memory_used_encode = _no_encode
elif "decoder.mid.block_1.mix_factor" in sd:
if "decoder.mid.block_1.mix_factor" in sd:
encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
decoder_config = encoder_config.copy()
decoder_config["video_kernel_size"] = [3, 1, 1]
@ -1710,16 +1692,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_target.params = {}
if len(clip_data) == 1:
te_model = detect_te_model(clip_data[0])
if clip_type == CLIPType.MINIMAX and "model.audio_decoder.projection.weight" in clip_data[0]:
tokenizer_data["tokenizer_json"] = clip_data[0].pop("tokenizer_json", None)
quant = comfy.utils.detect_layer_quantization(clip_data[0], "")
if quant is not None:
model_options = model_options.copy()
model_options["quantization_metadata"] = quant
clip_target.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0])
clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel
clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer
elif te_model == TEModel.CLIP_G:
if te_model == TEModel.CLIP_G:
if clip_type == CLIPType.STABLE_CASCADE:
clip_target.clip = sdxl_clip.StableCascadeClipModel
clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer

View File

@ -16,7 +16,6 @@ import comfy.text_encoders.genmo
import comfy.text_encoders.lt
import comfy.text_encoders.hunyuan_video
import comfy.text_encoders.minimax
import comfy.text_encoders.minimax_music
import comfy.text_encoders.cosmos
import comfy.text_encoders.lumina2
import comfy.text_encoders.wan
@ -2201,28 +2200,6 @@ class ACEStep15(supported_models_base.BASE):
return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
class MiniMaxMusic3(supported_models_base.BASE):
unet_config = {
"audio_model": "minimax_music3",
}
latent_format = comfy.latent_formats.MiniMaxMusic3
memory_usage_factor = 2.0
supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
sampling_settings = {"multiplier": 1.0}
def get_model(self, state_dict, prefix="", device=None):
return model_base.MiniMaxMusic3(self, device=device)
def model_type(self, state_dict, prefix=""):
return model_base.ModelType.FLOW
def clip_target(self, state_dict={}):
detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0])
target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
target.params["projection_config"] = detect
return target
class LongCatImage(supported_models_base.BASE):
unet_config = {
@ -2517,7 +2494,6 @@ models = [
ChromaRadiance,
ACEStep,
ACEStep15,
MiniMaxMusic3,
Omnigen2,
Boogu,
MageFlow,

View File

@ -6,16 +6,13 @@ import numpy as np
from tokenizers import Tokenizer
from dataclasses import dataclass
import math
import re
from comfy import sd1_clip
import comfy.model_management
import comfy.model_prefetch
import comfy.ops
import comfy.quant_ops
from comfy.ldm.modules.attention import optimized_attention_for_device
from comfy.rmsnorm import rms_norm
from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, FixedKV, _make_scaled_embedding
from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, _make_scaled_embedding
# Intentional minor divergences from transformers -reference implementation:
@ -112,28 +109,7 @@ class Gemma4_12B_Config(Gemma4Config):
suppress_tokens = [258883, 258882]
class RingKV(FixedKV):
# sliding-window ring: writes wrap at capacity, validity saturates
def prepare(self, num_tokens):
capacity = self.key.shape[2]
self.position.fill_(self.index % capacity)
self.seqlen.fill_(min(self.index + num_tokens, capacity))
def _fixed_kv_decode_mask(mask, cache, min_val):
capacity = cache.key.shape[2]
valid = min(cache.index + 1, capacity)
output = mask.new_full((*mask.shape[:-1], capacity), min_val)
if isinstance(cache, RingKV):
positions = torch.arange(cache.index + 1 - valid, cache.index + 1, device=mask.device) % capacity
output.index_copy_(-1, positions, mask[..., -valid:])
else:
output[..., :valid] = mask[..., :valid]
return output
# unfused RoPE as addcmul_ RoPE diverges from reference code (vision only; text
# layers use the kitchen split-half kernel, bitwise-equal to this with bf16 freqs)
# unfused RoPE as addcmul_ RoPE diverges from reference code
def _apply_rotary_pos_emb(x, freqs_cis):
cos, sin = freqs_cis[0], freqs_cis[1]
half = x.shape[-1] // 2
@ -164,23 +140,6 @@ class Gemma4Attention(nn.Module):
if config.k_norm == "gemma3":
self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps, device=device, dtype=dtype)
def _decode_attention(self, xq, cache, bias):
if bias is None:
# eager decode: slice the cache to the valid length (python-side index,
# no mask needed; a full ring is order-invariant under softmax)
n = min(cache.index + 1, cache.key.shape[2])
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
attention = optimized_attention_for_device(xq.device, mask=False, small_input=True)
return attention(xq, cache.key[:, :, :n], cache.value[:, :, :n], self.num_heads, skip_reshape=True, scale=1.0, **gqa_kwargs)
# graph capture: fixed-length masked attention over the full capacity, explicit
# math (SDPA leaves its fast path on broadcast-bias + GQA and costs ~0.5ms/layer)
batch_size = xq.shape[0]
groups = self.num_heads // self.num_kv_heads
q = xq.reshape(batch_size, self.num_kv_heads, groups, self.head_dim)
scores = q @ cache.key.transpose(-1, -2) + bias
probs = torch.softmax(scores.float(), dim=-1).to(xq.dtype)
return (probs @ cache.value).reshape(batch_size, 1, self.inner_size)
def forward(
self,
hidden_states: torch.Tensor,
@ -197,16 +156,10 @@ class Gemma4Attention(nn.Module):
if self.q_norm is not None:
xq = self.q_norm(xq)
if isinstance(shared_kv, FixedKV):
# decode on a KV-shared layer: attend the source layer's fixed cache
xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
output = self._decode_attention(xq, shared_kv, attention_mask)
return self.o_proj(output), None, None
if shared_kv is not None:
xk, xv = shared_kv
# Apply RoPE to Q only (K already has RoPE from source layer)
xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
xq = _apply_rotary_pos_emb(xq, freqs_cis)
present_key_value = None
shareable_kv = None
else:
@ -220,39 +173,11 @@ class Gemma4Attention(nn.Module):
xv = rms_norm(xv)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
xk = comfy.quant_ops.ck.apply_rope_split_half1(xk, freqs_cis)
xq = _apply_rotary_pos_emb(xq, freqs_cis)
xk = _apply_rotary_pos_emb(xk, freqs_cis)
present_key_value = None
fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None
if fixed_cache is not None:
if seq_length == 1:
# CUDA-graphable decode: write at the device-side ring/linear position
fixed_cache.key.index_copy_(2, fixed_cache.position, xk)
fixed_cache.value.index_copy_(2, fixed_cache.position, xv)
output = self._decode_attention(xq, fixed_cache, attention_mask)
return self.o_proj(output), fixed_cache, None
# prefill: attend the local sequence, persist the tail into the cache
capacity = fixed_cache.key.shape[2]
index = fixed_cache.index
if index + seq_length <= capacity:
fixed_cache.key[:, :, index:index + seq_length] = xk
fixed_cache.value[:, :, index:index + seq_length] = xv
if index > 0:
xk = fixed_cache.key[:, :, :index + seq_length]
xv = fixed_cache.value[:, :, :index + seq_length]
elif index == 0:
# prefill longer than the sliding ring: attend the full local K/V
# (per-query windows come from the prefill sliding mask), cache only
# the last `capacity` keys at their wrapped slots (position % capacity)
slots = torch.arange(seq_length - capacity, seq_length, device=xk.device) % capacity
fixed_cache.key.index_copy_(2, slots, xk[:, :, -capacity:])
fixed_cache.value.index_copy_(2, slots, xv[:, :, -capacity:])
else:
raise RuntimeError("gemma4: chunked prefill past the sliding window is not supported")
present_key_value = fixed_cache
elif past_key_value is not None:
if past_key_value is not None:
cumulative_len = 0
if len(past_key_value) > 0:
past_key, past_value, cumulative_len = past_key_value
@ -320,7 +245,6 @@ class TransformerBlockGemma4(nn.Module):
self.register_buffer("layer_scalar", torch.empty(1, device=device, dtype=dtype))
def forward(self, x, attention_mask=None, freqs_cis=None, past_key_value=None, per_layer_input=None, shared_kv=None):
output = x
sliding_window = None
if self.sliding_attention:
sliding_window = self.sliding_attention
@ -357,8 +281,7 @@ class TransformerBlockGemma4(nn.Module):
x = self.post_per_layer_input_norm(x)
x = residual + x
# in-place into the input buffer so CUDA-graph replays land in the static x
x = torch.mul(x, comfy.ops.cast_to_input(self.layer_scalar, x), out=output)
x = x * comfy.ops.cast_to_input(self.layer_scalar, x)
return x, present_key_value, shareable_kv
@ -367,9 +290,6 @@ class Gemma4Transformer(nn.Module):
def __init__(self, config, device=None, dtype=None, ops=None):
super().__init__()
self.config = config
self.fixed_kv = True
self.prefetch_dynamic_vbars = True
self.graph_dynamic_vbar_blocks = True
self.embed_tokens = _make_scaled_embedding(ops, config.vocab_size, config.hidden_size, config.hidden_size ** 0.5, device, dtype)
@ -378,19 +298,6 @@ class Gemma4Transformer(nn.Module):
for i in range(config.num_hidden_layers)
])
# KV-shared layers never run k_proj/v_proj/k_norm: their never-resolved vbar
# signatures would block layer graph capture, so prefetch only what executes
first_kv_shared = config.num_hidden_layers - config.num_kv_shared_layers if config.num_kv_shared_layers > 0 else config.num_hidden_layers
self._prefetch_units = []
for i, layer in enumerate(self.layers):
if i >= first_kv_shared:
dead = {layer.self_attn.k_proj, layer.self_attn.v_proj, layer.self_attn.k_norm}
self._prefetch_units.append([
m for m in layer.modules() if next(m.children(), None) is None and m not in dead
])
else:
self._prefetch_units.append(layer)
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, device=device, dtype=dtype) if config.final_norm else None
# Precompute RoPE inv_freq on CPU to match reference code's exact value
@ -404,9 +311,6 @@ class Gemma4Transformer(nn.Module):
sliding_inv = 1.0 / (config.rope_theta[1] ** (torch.arange(0, config.head_dim, 2).float() / config.head_dim))
self.register_buffer("_sliding_inv_freq", sliding_inv, persistent=False)
if config.suppress_tokens:
self.register_buffer("_suppress_tokens", torch.tensor(config.suppress_tokens, dtype=torch.long), persistent=False)
# Per-layer input mechanism
self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
if self.hidden_size_per_layer_input:
@ -418,26 +322,19 @@ class Gemma4Transformer(nn.Module):
self.hidden_size_per_layer_input, eps=config.rms_norm_eps,
device=device, dtype=dtype)
def get_dynamic_vram__units(self):
return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
def get_past_len(self, past_key_values):
for kv in past_key_values:
if isinstance(kv, FixedKV):
return kv.index
if len(kv) >= 3:
return kv[2]
return 0
def _freqs_from_inv(self, inv_freq, position_ids, device, dtype):
"""Compute per-pair 2x2 rotation matrices [B, 1, S, d/2, 2, 2] from stored inv_freq"""
"""Compute cos/sin from stored inv_freq"""
inv_exp = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(device)
pos_exp = position_ids[:, None, :].float()
freqs = (inv_exp @ pos_exp).transpose(1, 2)
cos, sin = freqs.cos(), freqs.sin()
mat = torch.stack((torch.stack((cos, -sin), dim=-1),
torch.stack((sin, cos), dim=-1)), dim=-2)
return mat.unsqueeze(1).to(dtype)
emb = torch.cat((freqs, freqs), dim=-1)
return emb.cos().unsqueeze(1).to(dtype), emb.sin().unsqueeze(1).to(dtype)
def compute_freqs_cis(self, position_ids, device, dtype=None):
global_freqs = self._freqs_from_inv(self._global_inv_freq, position_ids, device, dtype)
@ -504,71 +401,6 @@ class Gemma4Transformer(nn.Module):
first_kv_shared = self.config.num_hidden_layers - num_kv_shared if num_kv_shared > 0 else self.config.num_hidden_layers
shared_sliding_kv = None # KV from last non-shared sliding layer
shared_global_kv = None # KV from last non-shared global layer
share_source = {}
if num_kv_shared > 0:
for i in range(first_kv_shared):
share_source[bool(self.layers[i].sliding_attention)] = i
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(
list(self._prefetch_units), x.device,
{"prefetch_dynamic_vbars": self.prefetch_dynamic_vbars and past_key_values is not None})
fixed_kv = (past_key_values is not None and len(past_key_values) > 0
and isinstance(past_key_values[0], FixedKV))
decode = fixed_kv and seq_len == 1
# mirror the conditions under which prefetch_queue_pop can actually capture, so
# eager fallbacks keep the sliced decode path instead of the full-capacity one
enable_graph = (decode and mask is None and self.graph_dynamic_vbar_blocks
and prefetch_queue is not None
and hasattr(self.layers[0], "_v_block")
and not comfy.model_management.args.disable_cuda_graphs
and comfy.model_management.is_device_cuda(x.device))
decode_bias = None
decode_masks = None
if decode:
prepared = set()
for kv in past_key_values:
if isinstance(kv, FixedKV) and id(kv.position) not in prepared:
kv.prepare(seq_len)
prepared.add(id(kv.position))
if mask is not None:
decode_masks = {}
for kv in past_key_values:
if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
if enable_graph:
# static buffers + per-capacity attention biases: layer graphs replay against
# stable storage, refreshed eagerly each step
capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
state_key = (x.shape, x.dtype, x.device, tuple(t.shape for t in freqs_cis), capacities,
None if per_layer_inputs is None else per_layer_inputs.shape)
state = getattr(self, "_comfy_cross_step_state", None)
if state is None or state["key"] != state_key:
state = {"key": state_key,
"x": torch.empty_like(x),
"freqs_cis": [torch.empty_like(t) for t in freqs_cis],
"bias": {c: torch.empty((1, 1, 1, c), dtype=x.dtype, device=x.device) for c in capacities},
"per_layer": None if per_layer_inputs is None else torch.empty_like(per_layer_inputs),
"bias_valid": -1}
self._comfy_cross_step_state = state
comfy.model_management._register_cross_step(self)
state["x"].copy_(x)
for source, target in zip(freqs_cis, state["freqs_cis"]):
target.copy_(source)
x = state["x"]
freqs_cis = state["freqs_cis"]
if per_layer_inputs is not None:
state["per_layer"].copy_(per_layer_inputs)
per_layer_inputs = state["per_layer"]
valid = past_len + 1
for capacity, bias in state["bias"].items():
if state["bias_valid"] != past_len:
bias.fill_(min_val)
bias[..., :min(valid, capacity)] = 0
elif past_len < capacity:
bias[..., past_len:valid] = 0
state["bias_valid"] = valid
decode_bias = state["bias"]
intermediate = None
all_intermediate = None
@ -597,36 +429,12 @@ class Gemma4Transformer(nn.Module):
is_sliding = hasattr(layer, 'sliding_attention') and layer.sliding_attention
if i >= first_kv_shared and num_kv_shared > 0:
if decode:
layer_kwargs['shared_kv'] = past_key_values[share_source[bool(is_sliding)]]
else:
shared = shared_sliding_kv if is_sliding else shared_global_kv
if shared is not None:
layer_kwargs['shared_kv'] = shared
shared = shared_sliding_kv if is_sliding else shared_global_kv
if shared is not None:
layer_kwargs['shared_kv'] = shared
if enable_graph:
bias_cache = layer_kwargs.get('shared_kv', past_kv)
layer_mask = decode_bias[bias_cache.key.shape[2]]
elif decode:
bias_cache = layer_kwargs.get('shared_kv', past_kv)
layer_mask = None if decode_masks is None else decode_masks[id(bias_cache.position)]
else:
layer_mask = mask
x, current_kv, shareable_kv = layer(x=x, attention_mask=mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
result = []
def core():
nonlocal x
x, current_kv, shareable_kv = layer(x=x, attention_mask=layer_mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
result.append((current_kv, shareable_kv))
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
if result:
current_kv, shareable_kv = result[0]
else:
# graph replay: the cache already holds this step's write
current_kv, shareable_kv = past_kv, None
next_key_values.append(current_kv if current_kv is not None else ())
# Only track the last sliding/global before the sharing boundary
@ -639,14 +447,6 @@ class Gemma4Transformer(nn.Module):
if i == intermediate_output:
intermediate = x.clone()
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
if fixed_kv:
for kv in past_key_values:
if isinstance(kv, FixedKV):
kv.advance(seq_len)
if self.norm is not None:
x = self.norm(x)
@ -681,37 +481,14 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module):
if cap:
logits = cap * torch.tanh(logits / cap)
if self.model.config.suppress_tokens:
logits.index_fill_(-1, self.model._suppress_tokens, torch.finfo(logits.dtype).min)
logits[..., self.model.config.suppress_tokens] = torch.finfo(logits.dtype).min
return logits
def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
cfg = self.model.config
num_layers = cfg.num_hidden_layers
if not self.model.fixed_kv:
return [() for _ in range(num_layers)]
first_shared = num_layers - cfg.num_kv_shared_layers if cfg.num_kv_shared_layers > 0 else num_layers
# position/seqlen device tensors are shared per cache geometry and filled once per step
trackers = {}
caches = []
for i in range(num_layers):
if i >= first_shared:
caches.append(())
continue
sliding = cfg.sliding_attention[i % len(cfg.sliding_attention)] if cfg.sliding_attention else False
head_dim = cfg.head_dim if sliding else cfg.global_head_dim
k_eq_v = cfg.attention_k_eq_v and not sliding
kv_heads = cfg.num_global_key_value_heads if k_eq_v else cfg.num_key_value_heads
length = min(sliding, max_cache_len) if sliding else max_cache_len
cache_cls = RingKV if sliding else FixedKV
tracker = trackers.get((cache_cls, length))
if tracker is None:
tracker = (torch.empty((1,), device=device, dtype=torch.int64),
torch.empty((batch,), device=device, dtype=torch.int32))
trackers[(cache_cls, length)] = tracker
# zero-init: decode attends full capacity with masked tails, 0*0 stays finite
key = torch.zeros((batch, kv_heads, length, head_dim), device=device, dtype=execution_dtype)
caches.append(cache_cls(key, torch.zeros_like(key), 0, tracker[0], tracker[1]))
return caches
past_key_values = []
for _ in range(self.model.config.num_hidden_layers):
past_key_values.append(())
return past_key_values
def preprocess_embed(self, embed, device):
if embed["type"] == "image":
@ -1406,7 +1183,6 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po
class Gemma4_Tokenizer():
tokenizer_json_data = None
prime_empty_thought = False
def state_dict(self):
if self.tokenizer_json_data is not None:
@ -1557,8 +1333,8 @@ class Gemma4_Tokenizer():
num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
n_audio_tokens = self._audio_token_count(num_samples)
media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
# 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline.
model_open = "<|channel>thought\n<channel|>" if self.prime_empty_thought and not thinking else ""
# Non-thinking mode primes an empty thought channel so the model answers directly.
model_open = "" if thinking else "<|channel>thought\n<channel|>"
llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
@ -1625,13 +1401,11 @@ class Gemma4SDTokenizer(Gemma4_Tokenizer, sd1_clip.SDTokenizer):
def decode(self, token_ids, **kwargs):
text = super().decode(token_ids, skip_special_tokens=False)
# Only a close that ends a thought channel becomes </think>: generation primed with
# another channel leaves its opener in the prompt, so its close is not reasoning.
text = re.sub(r"<\|channel>thought\n(.*?)<channel\|>", r"<think>\n\1</think>", text, flags=re.DOTALL)
# Translate thinking channel markers to standard <think>/</think> tags
text = text.replace("<|channel>thought\n", "<think>\n")
text = text.replace("<channel|>", "</think>")
# Strip remaining special tokens
text = re.sub(r"<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?|<turn\|>", "", text)
text = text.replace("<eos>", "").strip()
text = text.replace("<turn|>", "").replace("<eos>", "").strip()
return text
@ -1644,7 +1418,6 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
"""Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
embedding_size = 3840
prime_empty_thought = True
def _extract_audio_features(self, waveform, sample_rate):
audio = self._resample_16k(waveform, sample_rate)
@ -1722,15 +1495,12 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
if dtype_llama is not None:
dtype = dtype_llama
super().__init__(device=device, dtype=dtype, name="gemma4", clip_model=clip_model, model_options=model_options)
def get_dynamic_vram__units(self):
return getattr(self, self.clip).transformer.model.get_dynamic_vram__units()
return Gemma4TEModel_
# Variants
def _make_variant(config_cls, prime_empty_thought=False):
def _make_variant(config_cls):
audio = config_cls.audio_config is not None
bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,)
class Variant(*bases):
@ -1740,8 +1510,8 @@ def _make_variant(config_cls, prime_empty_thought=False):
if audio:
self._init_audio(self.model.config, dtype, device, operations)
embedding_size = config_cls.hidden_size
if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought:
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought})
if embedding_size != Gemma4SDTokenizer.embedding_size:
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size})
class Tokenizer(Gemma4Tokenizer):
tokenizer_class = tok_cls
Variant.tokenizer = Tokenizer
@ -1751,7 +1521,7 @@ def _make_variant(config_cls, prime_empty_thought=False):
Gemma4_E4B = _make_variant(Gemma4Config)
Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True)
Gemma4_31B = _make_variant(Gemma4_31B_Config)
# Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).

View File

@ -5,33 +5,15 @@ from typing import Optional, Any, Tuple
import math
from tqdm import tqdm
import comfy.utils
import comfy_kitchen
from comfy.ldm.modules.attention import optimized_attention_for_device
import comfy.model_management
import comfy.model_prefetch
import comfy.ops
import comfy.ldm.common_dit
import comfy.clip_model
from . import qwen_vl
@dataclass
class FixedKV:
key: torch.Tensor
value: torch.Tensor
index: int
position: torch.Tensor
seqlen: torch.Tensor
def prepare(self, num_tokens):
self.position.fill_(self.index)
self.seqlen.fill_(self.index + num_tokens)
def advance(self, num_tokens):
self.index += num_tokens
@dataclass
class Llama2Config:
vocab_size: int = 128320
@ -267,9 +249,6 @@ class Qwen3_8BConfig:
rope_scale = None
final_norm: bool = True
lm_head: bool = True
fixed_kv: bool = False
merged_qkv: bool = False
merged_mlp: bool = False
stop_tokens = [151643, 151645]
@dataclass
@ -519,14 +498,9 @@ class Attention(nn.Module):
self.inner_size = self.num_heads * self.head_dim
ops = ops or nn
self.kv_size = self.num_kv_heads * self.head_dim
self.merged_qkv = getattr(config, "merged_qkv", False)
if self.merged_qkv:
self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype)
else:
self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.k_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.v_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.k_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype)
self.q_norm = None
@ -548,12 +522,9 @@ class Attention(nn.Module):
):
batch_size, seq_length, _ = hidden_states.shape
if self.merged_qkv:
xq, xk, xv = self.qkv_proj(hidden_states).split((self.inner_size, self.kv_size, self.kv_size), dim=-1)
else:
xq = self.q_proj(hidden_states)
xk = self.k_proj(hidden_states)
xv = self.v_proj(hidden_states)
xq = self.q_proj(hidden_states)
xk = self.k_proj(hidden_states)
xv = self.v_proj(hidden_states)
xq = xq.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
@ -566,29 +537,8 @@ class Attention(nn.Module):
xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None
if fixed_cache is not None:
xq = xq.transpose(1, 2)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
if seq_length == 1:
# CUDA-graphable decode path.
fixed_cache.key.index_copy_(1, fixed_cache.position, xk)
fixed_cache.value.index_copy_(1, fixed_cache.position, xv)
output = comfy_kitchen.flash_attention_decode(xq, fixed_cache.key, fixed_cache.value, fixed_cache.seqlen)
return self.o_proj(output.view(batch_size, seq_length, self.inner_size)), fixed_cache
fixed_cache.key[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xk)
fixed_cache.value[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xv)
xk = fixed_cache.key[:, :fixed_cache.index + seq_length]
xv = fixed_cache.value[:, :fixed_cache.index + seq_length]
xq = xq.transpose(1, 2)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
present_key_value = fixed_cache
if fixed_cache is None and past_key_value is not None:
present_key_value = None
if past_key_value is not None:
index = 0
num_tokens = xk.shape[2]
if len(past_key_value) > 0:
@ -619,27 +569,15 @@ class MLP(nn.Module):
def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None, intermediate_size=None):
super().__init__()
intermediate_size = intermediate_size or config.intermediate_size
self.merged_mlp = getattr(config, "merged_mlp", False)
if self.merged_mlp:
self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype)
else:
self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype)
if config.mlp_activation == "silu":
self.activation = torch.nn.functional.silu
self.merged_input_act = "swiglu"
elif config.mlp_activation == "gelu_pytorch_tanh":
self.activation = lambda a: torch.nn.functional.gelu(a, approximate="tanh")
self.merged_input_act = None
def forward(self, x):
if self.merged_mlp:
x = self.gate_up_proj(x)
if self.merged_input_act is not None:
return comfy.ops.linear_input_act(self.down_proj, x, self.merged_input_act)
gate, up = x.chunk(2, dim=-1)
return self.down_proj(self.activation(gate) * up)
return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))
class TransformerBlock(nn.Module):
@ -658,7 +596,6 @@ class TransformerBlock(nn.Module):
optimized_attention=None,
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
):
output = x
# Self Attention
residual = x
x = self.input_layernorm(x)
@ -675,7 +612,7 @@ class TransformerBlock(nn.Module):
residual = x
x = self.post_attention_layernorm(x)
x = self.mlp(x)
x = torch.add(residual, x, out=output)
x = residual + x
return x, present_key_value
@ -704,7 +641,6 @@ class TransformerBlockGemma2(nn.Module):
optimized_attention=None,
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
):
output = x
sliding_window = None
if self.transformer_type == 'gemma3':
if self.sliding_attention:
@ -740,7 +676,7 @@ class TransformerBlockGemma2(nn.Module):
x = self.pre_feedforward_layernorm(x)
x = self.mlp(x)
x = self.post_feedforward_layernorm(x)
x = torch.add(residual, x, out=output)
x = residual + x
return x, present_key_value
@ -752,14 +688,9 @@ def _make_scaled_embedding(ops, vocab_size, hidden_size, scale, device, dtype):
class Llama2_(nn.Module):
fixed_kv = False
graph_dynamic_vbar_blocks = False
def __init__(self, config, device=None, dtype=None, ops=None):
super().__init__()
self.config = config
self.fixed_kv = getattr(config, "fixed_kv", False)
self.graph_dynamic_vbar_blocks = False
self.vocab_size = config.vocab_size
if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
@ -782,27 +713,8 @@ class Llama2_(nn.Module):
if config.lm_head:
self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
def get_dynamic_vram__units(self):
return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
def get_past_len(self, past_key_values):
first = past_key_values[0]
return first.index if isinstance(first, FixedKV) else first[2]
def init_kv_cache(self, batch, capacity, device, dtype):
caches = []
fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device)
for _ in range(self.config.num_hidden_layers):
if fixed_kv:
key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype)
value = torch.empty_like(key)
position = torch.empty((1,), device=device, dtype=torch.int64)
seqlen = torch.empty((batch,), device=device, dtype=torch.int32)
caches.append(FixedKV(key, value, 0, position, seqlen))
else:
key = torch.empty((batch, self.config.num_key_value_heads, capacity, self.config.head_dim), device=device, dtype=dtype)
caches.append((key, torch.empty_like(key), 0))
return caches
return past_key_values[0][2]
def compute_freqs_cis(self, position_ids, device):
return precompute_freqs_cis(self.config.head_dim,
@ -844,33 +756,6 @@ class Llama2_(nn.Module):
optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV)
enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv and seq_len == 1 and mask is None
if enable_graph:
freqs_cis_groups = freqs_cis if isinstance(freqs_cis, list) else [freqs_cis]
cross_step_state_key = [(x.shape, x.stride(), x.dtype, x.device)]
for group in freqs_cis_groups:
for tensor in group:
cross_step_state_key.append((tensor.shape, tensor.stride(), tensor.dtype, tensor.device))
cross_step_state_key = tuple(cross_step_state_key)
cross_step_state = getattr(self, "_comfy_cross_step_state", None)
if cross_step_state is None or cross_step_state["key"] != cross_step_state_key:
static_freqs_cis = []
for group in freqs_cis_groups:
static_freqs_cis.append(tuple(torch.empty_like(tensor) for tensor in group))
if not isinstance(freqs_cis, list):
static_freqs_cis = static_freqs_cis[0]
cross_step_state = {"key": cross_step_state_key, "x": torch.empty_like(x), "freqs_cis": static_freqs_cis}
self._comfy_cross_step_state = cross_step_state
comfy.model_management._register_cross_step(self)
cross_step_state["x"].copy_(x)
static_freqs_cis_groups = cross_step_state["freqs_cis"] if isinstance(freqs_cis, list) else [cross_step_state["freqs_cis"]]
for source_group, target_group in zip(freqs_cis_groups, static_freqs_cis_groups):
for source, target in zip(source_group, target_group):
target.copy_(source)
x = cross_step_state["x"]
freqs_cis = cross_step_state["freqs_cis"]
intermediate = None
all_intermediate = None
only_layers = None
@ -884,8 +769,7 @@ class Llama2_(nn.Module):
elif intermediate_output < 0:
intermediate_output = len(self.layers) + intermediate_output
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": getattr(self, "prefetch_dynamic_vbars", False)})
next_key_values = list(past_key_values) if past_key_values is not None else []
next_key_values = []
for i, layer in enumerate(self.layers):
if all_intermediate is not None:
if only_layers is None or (i in only_layers):
@ -895,24 +779,16 @@ class Llama2_(nn.Module):
if past_key_values is not None:
past_kv = past_key_values[i] if len(past_key_values) > 0 else []
if fixed_kv:
past_kv.prepare(seq_len)
x, current_kv = layer(
x=x,
attention_mask=mask,
freqs_cis=freqs_cis,
optimized_attention=optimized_attention,
past_key_value=past_kv,
)
def core():
nonlocal x
x, current_kv = layer(
x=x,
attention_mask=mask,
freqs_cis=freqs_cis,
optimized_attention=optimized_attention,
past_key_value=past_kv,
)
if next_key_values:
next_key_values[i] = current_kv
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
if fixed_kv:
next_key_values[i].advance(seq_len)
if current_kv is not None:
next_key_values.append(current_kv)
# DeepStack: add per-layer visual features into the first len() decoder layers at image positions (Qwen3-VL)
if deepstack_embeds is not None and i < len(deepstack_embeds):
@ -921,9 +797,6 @@ class Llama2_(nn.Module):
if i == intermediate_output:
intermediate = x.clone()
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
if self.norm is not None:
x = self.norm(x)
@ -937,7 +810,7 @@ class Llama2_(nn.Module):
if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
intermediate = self.norm(intermediate)
if next_key_values:
if len(next_key_values) > 0:
return x, intermediate, next_key_values
else:
return x, intermediate
@ -1001,7 +874,12 @@ class BaseGenerate:
return torch.nn.functional.linear(input, weight, None)
def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
return self.model.init_kv_cache(batch, max_cache_len, device, execution_dtype)
model_config = self.model.config
past_key_values = []
for x in range(model_config.num_hidden_layers):
past_key_values.append((torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype),
torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0))
return past_key_values
def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None, embeds_info=None):
device = embeds.device

View File

@ -1,117 +0,0 @@
import torch
from tokenizers import Tokenizer
import comfy.ops
from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR
from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt
MODEL_CONFIG = {
"vocab_size": 200000,
"hidden_size": 4096,
"intermediate_size": 12288,
"num_hidden_layers": 36,
"num_attention_heads": 32,
"num_key_value_heads": 8,
"max_position_embeddings": 10240,
"rms_norm_eps": 1e-6,
"rope_theta": 1000000.0,
"head_dim": 128,
"audio_vocab_size": 1024,
"audio_num_codebooks": 8,
"decoder_num_heads": 16,
"decoder_intermediate_size": 6144,
"decoder_num_layers": 4,
}
def detect_merged_config(state_dict, prefix=""):
return {
"merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
"merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
"decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
"decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
}
class MiniMaxMusic3Tokenizer:
def __init__(self, embedding_directory=None, tokenizer_data={}):
tokenizer_json = tokenizer_data.get("tokenizer_json")
if tokenizer_json is None:
raise ValueError("MiniMax Music3 text encoder checkpoint is missing tokenizer_json")
if torch.is_tensor(tokenizer_json):
tokenizer_json = tokenizer_json.detach().cpu().numpy().tobytes()
self.tokenizer_json = tokenizer_json
self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8"))
for token, expected in SPECIAL_TOKEN_IDS.items():
if self.tokenizer.token_to_id(token) != expected:
raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}")
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
prompt = build_prompt(text, kwargs.get("lyrics", ""))
token_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids
return {
"minimax_music3": [[(token, 1.0) for token in token_ids]],
"seed": int(kwargs.get("seed", 0)),
"max_audio_frames": int(kwargs.get("max_audio_frames", MAX_AUDIO_FRAMES)),
"cfg_scale": float(kwargs.get("cfg_scale", CFG_SCALE)),
"top_k": int(kwargs.get("top_k", CFG_TOP_K)),
}
def state_dict(self):
return {"tokenizer_json": torch.frombuffer(bytearray(self.tokenizer_json), dtype=torch.uint8)}
def decode(self, token_ids, skip_special_tokens=True):
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
class MiniMaxMusic3TEModel(MiniMaxMusic3AR):
def __init__(self, device="cpu", dtype=None, model_options={}, projection_config=None):
dtype = torch.bfloat16
quant_config = model_options.get("quantization_metadata", None)
operations = model_options.get("custom_operations", None)
if operations is None:
operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast
super().__init__({**MODEL_CONFIG, **(projection_config or {})}, dtype, device, operations)
self.dtypes = {dtype}
self.execution_device = device
def set_clip_options(self, options):
self.execution_device = options.get("execution_device", self.execution_device)
def reset_clip_options(self):
pass
def get_dynamic_vram__units(self):
units, last_units = self.model.get_dynamic_vram__units()
if self.model.pruned_embedding:
last_units = [*last_units, self.model.embed_tokens_prefill]
return [(self.model.audio_decoder, self.model.audio_extra_embedding), *units], last_units
def encode_token_weights(self, token_weight_pairs):
token_ids = [token for token, _ in token_weight_pairs["minimax_music3"][0]]
input_ids = torch.tensor([token_ids], dtype=torch.long)
seed = token_weight_pairs["seed"]
max_audio_frames = token_weight_pairs["max_audio_frames"]
cfg_scale = token_weight_pairs["cfg_scale"]
top_k = token_weight_pairs["top_k"]
hidden = self.generate(input_ids, seed, max_audio_frames, self.execution_device, cfg_scale, top_k)
return hidden.unsqueeze(0), None, {}
def load_state_dict(self, state_dict, strict=True, assign=False):
if self.model.pruned_embedding is None:
self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict
if self.model.pruned_embedding:
del self.model.embed_tokens
else:
del self.model.embed_tokens_prefill, self.model.embed_tokens_audio
if self.model.pruned_lm_head is None:
self.model.pruned_lm_head = "model.lm_head_pruned.weight" in state_dict
if self.model.pruned_lm_head:
del self.model.lm_head
else:
del self.model.lm_head_pruned
return super().load_state_dict(state_dict, strict=strict, assign=assign)
def load_sd(self, state_dict):
return self.load_state_dict(state_dict, strict=False, assign=getattr(self, "can_assign_sd", False))

View File

@ -57,81 +57,6 @@ class BriaRemoveBackgroundRequest(BaseModel):
seed: int = Field(...)
class BriaGenFillRequest(BaseModel):
image: str = Field(...)
mask: str = Field(
...,
description="Binary mask defining the region to fill: white (255) pixels are generated, "
"black (0) pixels are preserved. Must have the same aspect ratio as the image.",
)
prompt: str = Field(...)
negative_prompt: str | None = Field(None)
refine_prompt: bool = Field(True)
seed: int = Field(...)
prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
visual_input_content_moderation: bool = Field(
False, description="If true, returns 422 on image or mask moderation failure."
)
visual_output_content_moderation: bool = Field(
False, description="If true, returns 422 on visual output moderation failure."
)
class BriaEraseRequest(BaseModel):
image: str = Field(...)
mask: str = Field(
...,
description="Binary mask defining the region to erase: white (255) pixels are removed, "
"black (0) pixels are preserved. Must have the same aspect ratio as the image.",
)
mask_type: str = Field("manual", description="'manual' for hand-drawn masks, 'automatic' for segmentation masks.")
visual_input_content_moderation: bool = Field(
False, description="If true, returns 422 on image or mask moderation failure."
)
visual_output_content_moderation: bool = Field(
False, description="If true, returns 422 on visual output moderation failure."
)
class BriaExpandRequest(BaseModel):
image: str = Field(...)
aspect_ratio: str | float | None = Field(
None,
description="Target ratio: a preset string (1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9) "
"or a float between 0.5 and 3.0. When set, the canvas/placement fields are ignored.",
)
canvas_size: list[int] | None = Field(None, description="Output canvas [width, height]; area up to 5000x5000.")
original_image_size: list[int] | None = Field(
None, description="Size [width, height] of the original image inside the canvas."
)
original_image_location: list[int] | None = Field(
None,
description="Top-left corner [x, y] of the original image inside the canvas; "
"values may fall outside the canvas, cropping the image.",
)
prompt: str | None = Field(None, description="If omitted, Bria auto-generates a prompt from the image.")
negative_prompt: str | None = Field(None)
seed: int = Field(...)
prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
visual_input_content_moderation: bool = Field(
False, description="If true, returns 422 on image moderation failure."
)
visual_output_content_moderation: bool = Field(
False, description="If true, returns 422 on visual output moderation failure."
)
class BriaIncreaseResolutionRequest(BaseModel):
image: str = Field(...)
desired_increase: int = Field(..., description="Resolution multiplier, 2 or 4.")
visual_input_content_moderation: bool = Field(
False, description="If true, returns 422 on image moderation failure."
)
visual_output_content_moderation: bool = Field(
False, description="If true, returns 422 on visual output moderation failure."
)
class BriaStatusResponse(BaseModel):
request_id: str = Field(...)
status_url: str = Field(...)
@ -147,26 +72,6 @@ class BriaRemoveBackgroundResponse(BaseModel):
result: BriaRemoveBackgroundResult | None = Field(None)
class BriaImageResult(BaseModel):
image_url: str = Field(...)
class BriaImageResultResponse(BaseModel):
status: str = Field(...)
result: BriaImageResult | None = Field(None)
class BriaExpandResult(BaseModel):
image_url: str = Field(...)
prompt: str | None = Field(None)
seed: int | None = Field(None)
class BriaExpandResponse(BaseModel):
status: str = Field(...)
result: BriaExpandResult | None = Field(None)
class BriaImageEditResult(BaseModel):
structured_prompt: str = Field(...)
image_url: str = Field(...)

View File

@ -116,7 +116,6 @@ class Seedance2TaskCreationRequest(BaseModel):
seed: int | None = Field(None, ge=0, le=2147483647)
watermark: bool | None = Field(None)
output_format: str | None = Field(None)
omni_reference_task_type: str | None = Field(None, description="One of: auto, reference, edit, extend.")
class TaskCreationResponse(BaseModel):

View File

@ -161,30 +161,12 @@ class Hailuo03TaskCreationRequest(BaseModel):
..., min_length=1
)
resolution: str = Field(...)
duration: int = Field(..., ge=4, le=15)
duration: int = Field(..., ge=5, le=15)
ratio: str | None = Field(None)
seed: int | None = Field(None, ge=0, le=4294967295)
aigc_watermark: bool | None = Field(None)
class Hailuo03ContextIRRequest(BaseModel):
model: str = Field(...)
content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
..., min_length=1
)
duration: int = Field(..., ge=4, le=15)
ratio: str | None = Field(None)
class Hailuo03RegenerationRequest(BaseModel):
model: str = Field(...)
content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
..., min_length=1
)
resolution: str = Field(...)
aigc_watermark: bool | None = Field(None)
class Hailuo03TaskCreationResponse(BaseModel):
task_id: str = Field(...)
@ -196,7 +178,6 @@ class Hailuo03TaskError(BaseModel):
class Hailuo03TaskContent(BaseModel):
url: str | None = Field(None)
prompt: str | None = Field(None)
class Hailuo03TaskUsage(BaseModel):

View File

@ -6,13 +6,7 @@ from typing_extensions import override
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api_nodes.apis.bria import (
BriaEditImageRequest,
BriaEraseRequest,
BriaExpandRequest,
BriaExpandResponse,
BriaGenFillRequest,
BriaImageEditResponse,
BriaImageResultResponse,
BriaIncreaseResolutionRequest,
BriaRemoveBackgroundRequest,
BriaRemoveBackgroundResponse,
BriaRemoveVideoBackgroundRequest,
@ -27,30 +21,13 @@ from comfy_api_nodes.util import (
convert_mask_to_image,
download_url_to_image_tensor,
download_url_to_video_output,
downscale_image_tensor_by_max_side,
get_image_dimensions,
poll_op,
sync_op,
upload_image_to_comfyapi,
upload_video_to_comfyapi,
validate_string,
validate_video_duration,
)
BRIA_MAX_OUTPUT_SIDE = 8192
BRIA_MIN_RATIO = 0.5
BRIA_MAX_RATIO = 3.0
BRIA_MIN_SHORT_SIDE = 224
def _upscaled_output_side(height: int, width: int, multiplier: int) -> int:
prescale = max(1.0, BRIA_MIN_SHORT_SIDE / min(height, width))
return round(max(height, width) * prescale * multiplier)
def _smallest_output_side(height: int, width: int, multiplier: int) -> int:
return round(max(height, width) / min(height, width) * BRIA_MIN_SHORT_SIDE * multiplier)
class BriaImageEditNode(IO.ComfyNode):
@ -266,503 +243,6 @@ class BriaRemoveImageBackground(IO.ComfyNode):
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
def _mask_to_binary_image(mask: Input.Image, action: str) -> torch.Tensor:
binary = (mask > 0.5).float()
if not binary.any():
raise ValueError(
f"The mask is empty, so there is nothing to {action}. Masks are binarized at 50%: "
f"areas painted at less than half opacity are ignored."
)
return convert_mask_to_image(binary)
def _validate_mask_aspect_ratio(image: Input.Image, mask: Input.Image) -> None:
ih, iw = image.shape[1], image.shape[2]
mh, mw = mask.shape[-2], mask.shape[-1]
if abs(iw * mh - ih * mw) > 0.01 * ih * mw:
raise ValueError(f"Mask must have the same aspect ratio as the image: image is {iw}x{ih}, mask is {mw}x{mh}.")
class BriaGenFill(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="BriaGenFill",
display_name="Bria Generative Fill",
category="partner/image/Bria",
description="Generate objects or scenery inside a masked region of an image using Bria.",
inputs=[
IO.Image.Input("image"),
IO.Mask.Input(
"mask",
tooltip="White areas are filled with generated content, black areas are preserved. "
"The mask is binarized before sending, so partially painted areas count as white. "
"Must have the same aspect ratio as the image.",
),
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Description of what to generate inside the masked region.",
),
IO.String.Input("negative_prompt", multiline=True, default=""),
IO.Boolean.Input(
"refine_prompt",
default=True,
tooltip="Automatically adjust the prompt for better results; "
"disable to use the prompt exactly as written.",
),
IO.Int.Input(
"seed",
default=42,
min=1,
max=2147483647,
step=1,
display_mode=IO.NumberDisplay.number,
control_after_generate=True,
),
IO.DynamicCombo.Input(
"moderation",
options=[
IO.DynamicCombo.Option("false", []),
IO.DynamicCombo.Option(
"true",
[
IO.Boolean.Input("prompt_content_moderation", default=False),
IO.Boolean.Input("visual_input_moderation", default=False),
IO.Boolean.Input("visual_output_moderation", default=False),
],
),
],
tooltip="Moderation settings",
),
],
outputs=[IO.Image.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.0429}""",
),
)
@classmethod
async def execute(
cls,
image: Input.Image,
mask: Input.Image,
prompt: str,
negative_prompt: str,
refine_prompt: bool,
seed: int,
moderation: InputModerationSettings,
) -> IO.NodeOutput:
validate_string(prompt, min_length=1)
_validate_mask_aspect_ratio(image, mask)
mask_image = _mask_to_binary_image(mask, "fill")
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/bria/v2/image/edit/gen_fill", method="POST"),
data=BriaGenFillRequest(
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
mask=await upload_image_to_comfyapi(
cls, mask_image, total_pixels=None, wait_label="Uploading mask"
),
prompt=prompt,
negative_prompt=negative_prompt if negative_prompt else None,
refine_prompt=refine_prompt,
seed=seed,
prompt_content_moderation=moderation.get("prompt_content_moderation", False),
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
),
response_model=BriaStatusResponse,
)
response = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
status_extractor=lambda r: r.status,
response_model=BriaImageResultResponse,
)
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
class BriaEraser(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="BriaEraser",
display_name="Bria Eraser",
category="partner/image/Bria",
description="Remove objects or areas outlined by a mask from an image using Bria.",
inputs=[
IO.Image.Input("image"),
IO.Mask.Input(
"mask",
tooltip="White areas are erased, black areas are preserved. "
"The mask is binarized before sending, so partially painted areas count as white. "
"Must have the same aspect ratio as the image.",
),
IO.Combo.Input(
"mask_type",
options=["manual", "automatic"],
tooltip="manual for hand-drawn or brush masks, "
"automatic for masks produced by segmentation models such as SAM.",
),
IO.DynamicCombo.Input(
"moderation",
options=[
IO.DynamicCombo.Option("false", []),
IO.DynamicCombo.Option(
"true",
[
IO.Boolean.Input("visual_input_moderation", default=False),
IO.Boolean.Input("visual_output_moderation", default=False),
],
),
],
tooltip="Moderation settings",
),
],
outputs=[IO.Image.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.0286}""",
),
)
@classmethod
async def execute(
cls,
image: Input.Image,
mask: Input.Image,
mask_type: str,
moderation: dict,
) -> IO.NodeOutput:
_validate_mask_aspect_ratio(image, mask)
mask_image = _mask_to_binary_image(mask, "erase")
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/bria/v2/image/edit/erase", method="POST"),
data=BriaEraseRequest(
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
mask=await upload_image_to_comfyapi(
cls, mask_image, total_pixels=None, wait_label="Uploading mask"
),
mask_type=mask_type,
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
),
response_model=BriaStatusResponse,
)
response = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
status_extractor=lambda r: r.status,
response_model=BriaImageResultResponse,
)
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
class BriaExpandImage(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="BriaExpandImage",
display_name="Bria Expand Image",
category="partner/image/Bria",
description="Expand an image beyond its borders with generated content using Bria.",
inputs=[
IO.Image.Input("image"),
IO.DynamicCombo.Input(
"expand_mode",
options=[
*[IO.DynamicCombo.Option(ratio, []) for ratio in
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"]],
IO.DynamicCombo.Option(
"custom_ratio",
[
IO.Int.Input(
"ratio_width",
default=21,
min=1,
max=100,
tooltip="Width side of the target ratio: 21 and 9 give 21:9.",
),
IO.Int.Input(
"ratio_height",
default=9,
min=1,
max=100,
tooltip="Height side of the target ratio: 21 and 9 give 21:9. "
f"Bria only accepts width/height between {BRIA_MIN_RATIO} and "
f"{BRIA_MAX_RATIO}, so anything taller than 1:2 needs the manual mode.",
),
],
),
IO.DynamicCombo.Option(
"manual",
[
IO.Int.Input("canvas_width", default=1000, min=64, max=5000),
IO.Int.Input("canvas_height", default=1000, min=64, max=5000),
IO.Int.Input(
"image_width",
default=500,
min=1,
max=5000,
tooltip="Width of the original image inside the canvas.",
),
IO.Int.Input(
"image_height",
default=500,
min=1,
max=5000,
tooltip="Height of the original image inside the canvas.",
),
IO.Int.Input(
"image_x",
default=250,
min=-5000,
max=5000,
tooltip="X position of the image's top-left corner inside the canvas; "
"may fall outside the canvas, cropping the image.",
),
IO.Int.Input(
"image_y",
default=250,
min=-5000,
max=5000,
tooltip="Y position of the image's top-left corner inside the canvas; "
"may fall outside the canvas, cropping the image.",
),
],
),
],
tooltip="Target shape of the expanded image: a preset aspect ratio, a custom ratio, "
"or manual placement of the original image on a canvas. "
"Manual is the only mode that can reach a canvas taller than 1:2.",
),
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Optional description of the expanded scene; "
"when empty, Bria generates one from the image.",
),
IO.String.Input("negative_prompt", multiline=True, default=""),
IO.Int.Input(
"seed",
default=42,
min=1,
max=2147483647,
step=1,
display_mode=IO.NumberDisplay.number,
control_after_generate=True,
),
IO.DynamicCombo.Input(
"moderation",
options=[
IO.DynamicCombo.Option("false", []),
IO.DynamicCombo.Option(
"true",
[
IO.Boolean.Input("prompt_content_moderation", default=False),
IO.Boolean.Input("visual_input_moderation", default=False),
IO.Boolean.Input("visual_output_moderation", default=False),
],
),
],
tooltip="Moderation settings",
),
],
outputs=[
IO.Image.Output(),
IO.String.Output(display_name="prompt", tooltip="The prompt used for the expansion; "
"auto-generated by Bria when the prompt input is empty."),
],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.0286}""",
),
)
@classmethod
async def execute(
cls,
image: Input.Image,
expand_mode: dict,
prompt: str,
negative_prompt: str,
seed: int,
moderation: InputModerationSettings,
) -> IO.NodeOutput:
mode = expand_mode["expand_mode"]
aspect_ratio = canvas_size = original_image_size = original_image_location = None
if mode == "manual":
canvas_size = [expand_mode["canvas_width"], expand_mode["canvas_height"]]
original_image_size = [expand_mode["image_width"], expand_mode["image_height"]]
original_image_location = [expand_mode["image_x"], expand_mode["image_y"]]
elif mode == "custom_ratio":
ratio_width, ratio_height = expand_mode["ratio_width"], expand_mode["ratio_height"]
aspect_ratio = ratio_width / ratio_height
if not BRIA_MIN_RATIO <= aspect_ratio <= BRIA_MAX_RATIO:
raise ValueError(
f"Bria accepts a width-to-height ratio between {BRIA_MIN_RATIO} and {BRIA_MAX_RATIO}: "
f"{ratio_width}:{ratio_height} is {aspect_ratio:.4f}. "
f"Use the manual expand mode to reach a canvas of any shape."
)
else:
aspect_ratio = mode
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/bria/v2/image/edit/expand", method="POST"),
data=BriaExpandRequest(
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
aspect_ratio=aspect_ratio,
canvas_size=canvas_size,
original_image_size=original_image_size,
original_image_location=original_image_location,
prompt=prompt if prompt else None,
negative_prompt=negative_prompt if negative_prompt else None,
seed=seed,
prompt_content_moderation=moderation.get("prompt_content_moderation", False),
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
),
response_model=BriaStatusResponse,
)
response = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
status_extractor=lambda r: r.status,
response_model=BriaExpandResponse,
)
return IO.NodeOutput(
await download_url_to_image_tensor(response.result.image_url),
response.result.prompt or "",
)
class BriaIncreaseResolution(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="BriaIncreaseResolution",
display_name="Bria Increase Resolution",
category="partner/image/Bria",
description="Upscale an image by 2x or 4x using Bria, preserving the original content.",
inputs=[
IO.Image.Input("image"),
IO.Combo.Input(
"desired_increase",
options=["2", "4"],
tooltip="Resolution multiplier. The output must fit within 8192 pixels on each side.",
),
IO.Boolean.Input(
"auto_downscale",
default=False,
tooltip="Automatically lower the multiplier, and downscale the input image if that is "
"still not enough, when the output would exceed the limit.",
),
IO.DynamicCombo.Input(
"moderation",
options=[
IO.DynamicCombo.Option("false", []),
IO.DynamicCombo.Option(
"true",
[
IO.Boolean.Input("visual_input_moderation", default=False),
IO.Boolean.Input("visual_output_moderation", default=False),
],
),
],
tooltip="Moderation settings",
),
],
outputs=[IO.Image.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.0286}""",
),
)
@classmethod
async def execute(
cls,
image: Input.Image,
desired_increase: str,
auto_downscale: bool,
moderation: dict,
) -> IO.NodeOutput:
multiplier = int(desired_increase)
height, width = get_image_dimensions(image)
if _upscaled_output_side(height, width, multiplier) > BRIA_MAX_OUTPUT_SIDE:
candidates = [c for c in (4, 2) if c <= multiplier]
if not auto_downscale:
predicted = _upscaled_output_side(height, width, multiplier)
raise ValueError(
f"Bria can upscale up to a maximum output dimension of {BRIA_MAX_OUTPUT_SIDE} pixels: "
f"input is {width}x{height}, x{multiplier} would be {predicted} pixels on the long side. "
f"Enable auto_downscale, or use a smaller input image or a lower multiplier."
)
fitted = next(
(c for c in candidates if _upscaled_output_side(height, width, c) <= BRIA_MAX_OUTPUT_SIDE), None
)
if fitted is not None:
multiplier = fitted
else:
shrinkable = next((c for c in sorted(candidates) if _smallest_output_side(height, width, c)
<= BRIA_MAX_OUTPUT_SIDE), None)
if shrinkable is None:
raise ValueError(
f"This image cannot be upscaled by Bria at any multiplier: it is {width}x{height}, and "
f"Bria first enlarges the short side to {BRIA_MIN_SHORT_SIDE} pixels, which pushes the "
f"long side past the {BRIA_MAX_OUTPUT_SIDE} pixel limit. Crop it to a squarer shape first."
)
multiplier = shrinkable
image = downscale_image_tensor_by_max_side(image, max_side=BRIA_MAX_OUTPUT_SIDE // multiplier)
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/bria/v2/image/edit/increase_resolution", method="POST"),
data=BriaIncreaseResolutionRequest(
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
desired_increase=multiplier,
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
),
response_model=BriaStatusResponse,
)
response = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
status_extractor=lambda r: r.status,
response_model=BriaImageResultResponse,
)
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
class BriaRemoveVideoBackground(IO.ComfyNode):
@classmethod
@ -1092,10 +572,6 @@ class BriaExtension(ComfyExtension):
return [
BriaImageEditNode,
BriaRemoveImageBackground,
BriaGenFill,
BriaEraser,
BriaExpandImage,
BriaIncreaseResolution,
BriaRemoveVideoBackground,
BriaVideoGreenScreen,
BriaVideoReplaceBackground,

View File

@ -56,8 +56,6 @@ from comfy_api_nodes.util import (
ApiEndpoint,
audio_bytes_to_audio_input,
audio_input_to_mp3,
bytesio_to_image_tensor,
download_url_as_bytesio,
download_url_to_image_tensor,
download_url_to_video_output,
downscale_image_tensor_by_max_side,
@ -1317,9 +1315,7 @@ class ByteDanceSeedreamLayerSeparationNode(IO.ComfyNode):
left, top, rect_w, rect_h = spec["left"], spec["top"], spec["rect_w"], spec["rect_h"]
async with semaphore:
try:
# the layer math below needs the alpha channel, and ByteDance encodes
# alpha-less images as plain RGB (the base plate is one), so force RGBA
rgba = bytesio_to_image_tensor(await download_url_as_bytesio(str(item["url"])), mode="RGBA")[0]
rgba = (await download_url_to_image_tensor(str(item["url"])))[0]
except ProcessingInterrupted:
raise
except Exception as exc:
@ -2073,7 +2069,7 @@ def _seedance2_text_inputs(resolutions: list[str], default_ratio: str = "16:9"):
]
def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool = False, with_task_type: bool = False):
def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool = False):
return [
IO.String.Input(
"prompt",
@ -2128,29 +2124,6 @@ def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool =
if with_video_editing
else []
),
*(
[
IO.Combo.Input(
"task_type",
options=["auto", "reference", "edit", "extend"],
default="auto",
tooltip="What to do with the reference media. Every value except auto is "
"validated when the task is submitted, so mismatched settings fail before "
"generation starts. auto: the model infers the task from the prompt and "
"inputs, and settings that conflict with its reading fail only after "
"generation has started. reference: generate a new video guided by the "
"reference images, videos, and audio. edit: change a connected reference "
"video (add, remove, replace); the output keeps the source clip's own length "
"and aspect ratio, and the duration and ratio widgets are ignored. extend: "
"continue a connected reference video forward or backward; the prompt should "
"say 'extend forward', 'extend backward', or 'continue', the aspect ratio "
"follows the source clip, and the output contains only the newly generated "
"segment of the duration you set, not the source clip.",
)
]
if with_task_type
else []
),
IO.Combo.Input(
"output_format",
options=["mp4"],
@ -2160,9 +2133,9 @@ def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool =
]
def _seedance25_reference_inputs(with_video_editing: bool = False, with_task_type: bool = False):
def _seedance25_reference_inputs():
return [
*_seedance25_text_inputs(with_video_editing=with_video_editing, with_task_type=with_task_type),
*_seedance25_text_inputs(with_video_editing=True),
IO.Autogrow.Input(
"reference_images",
template=IO.Autogrow.TemplateNames(
@ -2223,23 +2196,17 @@ def _seedance2_build_request(
watermark: bool,
ratio: str,
) -> Seedance2TaskCreationRequest:
task_type = model.get("task_type", "auto")
duration = model["duration"]
if model.get("video_editing") or task_type == "edit":
ratio, duration = "adaptive", -1
elif task_type == "extend":
ratio = "adaptive"
video_editing = bool(model.get("video_editing"))
return Seedance2TaskCreationRequest(
model=model_id,
content=content,
generate_audio=model["generate_audio"],
resolution=model["resolution"],
ratio=ratio,
duration=duration,
ratio="adaptive" if video_editing else ratio,
duration=-1 if video_editing else model["duration"],
seed=seed,
watermark=watermark,
output_format=model.get("output_format"),
omni_reference_task_type=None if task_type == "auto" else task_type,
)
@ -2249,7 +2216,7 @@ _SEEDANCE2_PRICE_EXPR_TEMPLATE = """
$res := $lookup(widgets, "model.resolution");
$ratio := $lookup(widgets, "model.ratio");
$dur := $lookup(widgets, "model.duration");
$auto := __IS_EDIT__;
$auto := $lookup(widgets, "model.video_editing") = true;
$hasVideo := __HAS_VIDEO__;
$ready := $type($m) = "string" and $type($res) = "string" and ($auto or $type($dur) = "number");
$ready ? (
@ -2294,7 +2261,6 @@ _SEEDANCE2_PRICE_EXPR_TEMPLATE = """
_SEEDANCE_AUDIO_POLICY_CODE = "OutputAudioSensitiveContentDetected.PolicyViolation"
_SEEDANCE_TASK_TYPE_CONSTRAINT_CODE = "InvalidParameter.TaskTypeConstraint"
_SEEDANCE_TASK_TYPE_MISMATCH_CODE = "InvalidParameter.TaskTypeMismatch"
async def _seedance2_poll_video_task(
@ -2303,7 +2269,6 @@ async def _seedance2_poll_video_task(
model_id: str,
resolution: str,
has_video_input: bool,
task_type: str | None = None,
) -> TaskStatusResponse:
try:
return await poll_op(
@ -2324,48 +2289,19 @@ async def _seedance2_poll_video_task(
"to get a silent video, or adjust the prompt and try again."
) from exc
if _SEEDANCE_TASK_TYPE_CONSTRAINT_CODE in str(exc):
if task_type is None:
raise ValueError(
"Seedance read this prompt as editing the reference video, and an edit always "
"takes its duration and aspect ratio from that video. Enable video_editing on "
"this node and run again, or reword the prompt so it describes a new video "
"rather than a change to the reference one."
) from exc
if task_type == "edit":
raise ValueError(
"The request does not satisfy the 'edit' constraints: the clip being edited "
"must be 4 to 30 seconds long."
) from exc
if task_type == "extend":
raise ValueError(
"The request does not satisfy the 'extend' constraints: the clip being "
"extended must be 1.9 to 30 seconds long."
) from exc
raise ValueError(
"Seedance decided from the prompt that this task's duration or aspect ratio "
"must come from the reference video, and the current settings conflict with "
"that. Set task_type to the task you mean ('edit' or 'extend') and run again, "
"or reword the prompt so it describes a new video rather than a change to the "
"reference one."
) from exc
if _SEEDANCE_TASK_TYPE_MISMATCH_CODE in str(exc):
raise ValueError(
f"Seedance read this prompt as a different task than the selected task_type "
f"'{task_type}'. Reword the prompt so it matches: an extend prompt should say "
"'extend forward', 'extend backward', or 'continue'; an edit prompt should use "
"words like add, remove, replace, or change. Or set task_type to auto."
"Seedance read this prompt as editing the reference video, and an edit always "
"takes its duration and aspect ratio from that video. Enable video_editing on "
"this node and run again, or reword the prompt so it describes a new video "
"rather than a change to the reference one."
) from exc
raise
def _seedance2_price_badge(with_reference_videos: bool, legacy_video_editing: bool = False) -> IO.PriceBadge:
def _seedance2_price_badge(with_reference_videos: bool) -> IO.PriceBadge:
widgets = ["model", "model.resolution", "model.ratio", "model.duration"]
if legacy_video_editing:
is_edit = '$lookup(widgets, "model.video_editing") = true'
else:
is_edit = '$lookup(widgets, "model.task_type") = "edit"'
if with_reference_videos:
widgets.append("model.video_editing" if legacy_video_editing else "model.task_type")
widgets.append("model.video_editing")
has_video = (
'$exists(inputGroups) and $lookup(inputGroups, "model.reference_videos") > 0'
if with_reference_videos
@ -2376,7 +2312,7 @@ def _seedance2_price_badge(with_reference_videos: bool, legacy_video_editing: bo
widgets=widgets,
input_groups=["model.reference_videos"] if with_reference_videos else [],
),
expr=_SEEDANCE2_PRICE_EXPR_TEMPLATE.replace("__HAS_VIDEO__", has_video).replace("__IS_EDIT__", is_edit),
expr=_SEEDANCE2_PRICE_EXPR_TEMPLATE.replace("__HAS_VIDEO__", has_video),
)
@ -2726,12 +2662,12 @@ def _seedance2_reference_inputs(resolutions: list[str], default_ratio: str = "16
]
class ByteDance2ReferenceNodeV2(IO.ComfyNode):
class ByteDance2ReferenceNode(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="ByteDance2ReferenceNodeV2",
node_id="ByteDance2ReferenceNode",
display_name="ByteDance Seedance 2.5 Reference to Video",
category="partner/video/ByteDance",
description="Generate, edit, or extend video using Seedance 2.5 or 2.0 with reference "
@ -2740,7 +2676,7 @@ class ByteDance2ReferenceNodeV2(IO.ComfyNode):
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs(with_task_type=True)),
IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs()),
IO.DynamicCombo.Option(
"Seedance 2.0",
_seedance2_reference_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
@ -2825,13 +2761,6 @@ class ByteDance2ReferenceNodeV2(IO.ComfyNode):
f"(videos={len(reference_videos)}, video assets={len(reference_video_assets)}). "
f"Maximum is {limits['max_videos']}."
)
task_type = model.get("task_type")
if task_type in ("edit", "extend") and total_videos == 0:
raise ValueError(
f"A '{task_type}' task needs at least one reference video. Connect the video "
f"you want to {'change' if task_type == 'edit' else 'continue'}, or set "
"task_type to 'reference' to generate a new video from the references you have."
)
total_audios = len(reference_audios) + len(reference_audio_assets)
if total_audios > limits["max_audios"]:
raise ValueError(
@ -2964,78 +2893,11 @@ class ByteDance2ReferenceNodeV2(IO.ComfyNode):
response_model=TaskCreationResponse,
)
response = await _seedance2_poll_video_task(
cls,
initial_response.id,
model_id,
model["resolution"],
has_video_input=has_video_input,
task_type=task_type,
cls, initial_response.id, model_id, model["resolution"], has_video_input=has_video_input
)
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
class ByteDance2ReferenceNode(ByteDance2ReferenceNodeV2):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="ByteDance2ReferenceNode",
display_name="ByteDance Seedance 2.5 Reference to Video (Legacy)",
category="partner/video/ByteDance",
description="Generate, edit, or extend video using Seedance 2.5 or 2.0 with reference "
"images, videos, and audio. Supports multimodal reference, video editing, and video extension.",
inputs=[
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs(with_video_editing=True)),
IO.DynamicCombo.Option(
"Seedance 2.0",
_seedance2_reference_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
),
IO.DynamicCombo.Option(
"Seedance 2.0 Fast",
_seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
),
IO.DynamicCombo.Option(
"Seedance 2.0 Mini",
_seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
),
],
tooltip=SEEDANCE_MODEL_TOOLTIP,
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=2147483647,
step=1,
display_mode=IO.NumberDisplay.number,
control_after_generate=True,
tooltip="Seed controls whether the node should re-run; "
"results are non-deterministic regardless of seed.",
),
IO.Boolean.Input(
"watermark",
default=False,
tooltip="Whether to add a watermark to the video.",
advanced=True,
),
],
outputs=[
IO.Video.Output(),
],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
is_deprecated=True,
price_badge=_seedance2_price_badge(with_reference_videos=True, legacy_video_editing=True),
)
async def process_video_task(
cls: type[IO.ComfyNode],
payload: Text2VideoTaskCreationRequest | Image2VideoTaskCreationRequest,
@ -3551,7 +3413,6 @@ class ByteDanceExtension(ComfyExtension):
ByteDance2TextToVideoNode,
ByteDance2FirstLastFrameNode,
ByteDance2ReferenceNode,
ByteDance2ReferenceNodeV2,
ByteDanceCreateImageAsset,
ByteDanceCreateVideoAsset,
ByteDanceSeedAudioNode,

View File

@ -43,7 +43,6 @@ from comfy_api_nodes.util import (
download_url_to_image_tensor,
download_url_to_video_output,
get_number_of_images,
pad_images_to_common_channels,
sync_op,
tensor_to_base64_string,
upload_audio_to_comfyapi,
@ -234,8 +233,8 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
"Try rephrasing your prompt or changing the response modality to 'IMAGE+TEXT' "
"to see the model's reasoning."
)
return torch.zeros((1, 1024, 1024, 3))
return torch.cat(pad_images_to_common_channels(image_tensors), dim=0)
return torch.zeros((1, 1024, 1024, 4))
return torch.cat(image_tensors, dim=0)
def get_text_from_interaction(interaction: GeminiInteraction) -> str:

View File

@ -3,14 +3,12 @@ from typing import Optional
import torch
from typing_extensions import override
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api.latest import IO, ComfyExtension
from comfy_api_nodes.apis.minimax import (
Hailuo03AudioContent,
Hailuo03AudioContentUrl,
Hailuo03ContextIRRequest,
Hailuo03ImageContent,
Hailuo03ImageContentUrl,
Hailuo03RegenerationRequest,
Hailuo03TaskCreationRequest,
Hailuo03TaskCreationResponse,
Hailuo03TaskQueryResponse,
@ -458,9 +456,6 @@ HAILUO_03_QUERY_ENDPOINT = "/proxy/minimax/v2/query/video_generation" # + /{tas
HAILUO_03_MODELS = {"MiniMax H3": "MiniMax-H3"}
HAILUO_03_FAILED_STATUSES = ["failed", "cancelled", "expired"]
HAILUO_03_CONTEXT_IR_ENDPOINT = "/proxy/minimax/v2/h3_context_ir"
HAILUO_03_REGENERATION_ENDPOINT = "/proxy/minimax/v2/video_regeneration"
def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = True):
inputs = [
@ -492,10 +487,10 @@ def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = Tr
IO.Int.Input(
"duration",
default=5,
min=4,
min=5,
max=15,
step=1,
tooltip="Duration of the output video in seconds (4-15).",
tooltip="Duration of the output video in seconds (5-15).",
display_mode=IO.NumberDisplay.slider,
)
)
@ -944,592 +939,6 @@ class MinimaxHailuo03ReferenceNode(IO.ComfyNode):
)
class MinimaxHailuo03ContextIRNode(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="MinimaxHailuo03ContextIRNode",
display_name="MiniMax H3 Context IR (Prompt Enhancer)",
category="partner/video/MiniMax",
description="Analyze text and media context with MiniMax H3 Context IR and produce an enhanced, "
"structured video prompt. Feed the output into the prompt of a MiniMax H3 video node and attach "
"the same media there in the same order, because the enhanced prompt refers to the attached "
"media by position.",
inputs=[
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"MiniMax H3",
[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Description of the video you intend to generate.",
),
IO.Int.Input(
"duration",
default=5,
min=4,
max=15,
step=1,
tooltip="Duration of the video you intend to generate, in seconds (4-15).",
display_mode=IO.NumberDisplay.slider,
),
IO.Combo.Input(
"ratio",
options=["adaptive", "16:9", "4:3", "1:1", "3:4", "9:16", "21:9"],
default="adaptive",
tooltip="Aspect ratio of the video you intend to generate. 'adaptive' "
"requires at least one image, video, or audio input.",
),
IO.Autogrow.Input(
"reference_images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("reference_image"),
names=[
"image_1",
"image_2",
"image_3",
"image_4",
"image_5",
"image_6",
"image_7",
"image_8",
"image_9",
],
min=0,
),
tooltip="Subject or style reference images, referred to in the prompt "
"as 'Image 1'..'Image 9' in connection order. Up to 9 images.",
),
IO.Autogrow.Input(
"reference_videos",
template=IO.Autogrow.TemplateNames(
IO.Video.Input("reference_video"),
names=["video_1", "video_2", "video_3"],
min=0,
),
tooltip="Motion or scene reference videos, referred to in the prompt "
"as 'Video 1'..'Video 3' in connection order. Up to 3 videos, "
"2-15 seconds each, 15 seconds in total.",
),
IO.Autogrow.Input(
"reference_audios",
template=IO.Autogrow.TemplateNames(
IO.Audio.Input("reference_audio"),
names=["audio_1", "audio_2", "audio_3"],
min=0,
),
tooltip="Audio references, referred to in the prompt as "
"'Audio 1'..'Audio 3' in connection order. Up to 3 clips, "
"2-15 seconds each, 15 seconds in total. Cannot be used without "
"a reference image or video.",
),
],
)
],
tooltip="Model to use for prompt enhancement.",
),
IO.Image.Input(
"first_frame",
tooltip="First frame of the video you intend to generate. Cannot be combined with "
"reference media.",
optional=True,
),
IO.Image.Input(
"last_frame",
tooltip="Last frame of the video you intend to generate. Cannot be combined with "
"reference media.",
optional=True,
),
],
outputs=[
IO.String.Output(),
],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(
inputs=["first_frame", "last_frame"],
input_groups=["model.reference_images", "model.reference_videos", "model.reference_audios"],
),
expr="""
(
$imgsRaw := $lookup(inputGroups, "model.reference_images");
$imgs := $imgsRaw ? $imgsRaw : 0;
$vidsRaw := $lookup(inputGroups, "model.reference_videos");
$vids := $vidsRaw ? $vidsRaw : 0;
$audsRaw := $lookup(inputGroups, "model.reference_audios");
$auds := $audsRaw ? $audsRaw : 0;
$frames := (inputs.first_frame.connected ? 1 : 0) + (inputs.last_frame.connected ? 1 : 0);
($imgs + $vids + $auds) > 0
? {"type": "range_usd", "min_usd": 0.06, "max_usd": 0.11, "format": {"approximate": true}}
: $frames > 0
? {"type": "usd", "usd": 0.05, "format": {"approximate": true}}
: {"type": "usd", "usd": 0.02, "format": {"approximate": true}}
)
""",
),
)
@classmethod
async def execute(
cls,
model: dict,
first_frame: torch.Tensor | None = None,
last_frame: torch.Tensor | None = None,
) -> IO.NodeOutput:
validate_string(model["prompt"], strip_whitespace=True, min_length=1)
reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
has_frames = first_frame is not None or last_frame is not None
has_references = bool(reference_images) or bool(reference_videos) or bool(reference_audios)
if has_frames and has_references:
raise ValueError(
"First/last frame and reference media are mutually exclusive. Use frames for an "
"image-to-video prompt, or reference media for a reference-to-video prompt."
)
if reference_audios and not reference_images and not reference_videos:
raise ValueError("Reference audio cannot be used without a reference image or video.")
if not has_frames and not has_references and model["ratio"] == "adaptive":
raise ValueError(
"Ratio 'adaptive' is not supported for text-only requests; select an explicit aspect ratio."
)
for frame in (first_frame, last_frame):
if frame is not None:
validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
validate_image_dimensions(frame, min_width=256, min_height=256)
for image in reference_images.values():
validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
validate_image_dimensions(image, min_width=256, min_height=256)
total_video_duration = 0.0
for i, video in enumerate(reference_videos.values(), 1):
try:
fps = float(video.get_frame_rate())
except Exception:
fps = 0.0
if fps and not (23.9 <= fps <= 60.5):
raise ValueError(f"Reference video {i} is {fps:.2f} FPS. Supported range is 23.976-60 FPS.")
try:
dur = video.get_duration()
except Exception:
continue
if dur < 1.8:
raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
total_video_duration += dur
if total_video_duration > 15.1:
raise ValueError(
f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
)
total_audio_duration = 0.0
for i, audio in enumerate(reference_audios.values(), 1):
dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
if dur < 1.8:
raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
total_audio_duration += dur
if total_audio_duration > 15.1:
raise ValueError(
f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
)
content: list = [Hailuo03TextContent(text=model["prompt"])]
if first_frame is not None:
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, first_frame, max_images=1, wait_label="Uploading first frame"
)
)[0],
),
role="first_frame",
)
)
if last_frame is not None:
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, last_frame, max_images=1, wait_label="Uploading last frame"
)
)[0],
),
role="last_frame",
)
)
for i, image in enumerate(reference_images.values(), 1):
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, image, max_images=1, wait_label=f"Uploading image {i}"
)
)[0],
),
role="reference_image",
)
)
for i, video in enumerate(reference_videos.values(), 1):
content.append(
Hailuo03VideoContent(
video_url=Hailuo03VideoContentUrl(
url=await upload_video_to_comfyapi(cls, video, wait_label=f"Uploading video {i}"),
),
)
)
for audio in reference_audios.values():
content.append(
Hailuo03AudioContent(
audio_url=Hailuo03AudioContentUrl(
url=await upload_audio_to_comfyapi(
cls,
audio,
container_format="mp3",
codec_name="libmp3lame",
mime_type="audio/mpeg",
),
),
)
)
response = await sync_op(
cls,
ApiEndpoint(path=HAILUO_03_CONTEXT_IR_ENDPOINT, method="POST"),
response_model=Hailuo03TaskCreationResponse,
data=Hailuo03ContextIRRequest(
model=HAILUO_03_MODELS[model["model"]],
content=content,
duration=model["duration"],
ratio=None if model["ratio"] == "adaptive" else model["ratio"],
),
)
task_result = await poll_op(
cls,
ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
response_model=Hailuo03TaskQueryResponse,
status_extractor=lambda r: r.task.status,
failed_statuses=HAILUO_03_FAILED_STATUSES,
poll_interval=5,
)
prompt = task_result.task.content.prompt if task_result.task.content else None
if not prompt:
raise Exception(f"No enhanced prompt in the response: {task_result.model_dump()}")
return IO.NodeOutput(prompt)
class MinimaxHailuo03RegenerateNode(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="MinimaxHailuo03RegenerateNode",
display_name="MiniMax H3 Regenerate to 2K",
category="partner/video/MiniMax",
description="Re-render a MiniMax H3 768P output at 2K resolution. Connect the unmodified 768P "
"video and the exact prompt used to generate it; if the original generation used first/last "
"frames or reference media, attach the same inputs.",
inputs=[
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"MiniMax H3",
[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="The exact prompt used to generate the source video.",
),
IO.Combo.Input(
"resolution",
options=["2K"],
tooltip="Resolution to re-render the source video at.",
),
IO.Autogrow.Input(
"reference_images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("reference_image"),
names=[
"image_1",
"image_2",
"image_3",
"image_4",
"image_5",
"image_6",
"image_7",
"image_8",
"image_9",
],
min=0,
),
tooltip="Reference images from the original generation, in the same "
"order. Up to 9 images.",
),
IO.Autogrow.Input(
"reference_videos",
template=IO.Autogrow.TemplateNames(
IO.Video.Input("reference_video"),
names=["video_1", "video_2", "video_3"],
min=0,
),
tooltip="Reference videos from the original generation, in the same "
"order. Up to 3 videos, 2-15 seconds each, 15 seconds in total.",
),
IO.Autogrow.Input(
"reference_audios",
template=IO.Autogrow.TemplateNames(
IO.Audio.Input("reference_audio"),
names=["audio_1", "audio_2", "audio_3"],
min=0,
),
tooltip="Audio references from the original generation, in the same "
"order. Up to 3 clips, 2-15 seconds each, 15 seconds in total. "
"Cannot be used without a reference image or video.",
),
],
)
],
tooltip="Model to use for video regeneration.",
),
IO.Video.Input(
"video",
tooltip="The MiniMax H3 768P output video to re-render. Connect the unmodified output "
"of a MiniMax H3 video node (24 FPS, 4-15 seconds). 2K outputs cannot be used.",
),
IO.Image.Input(
"first_frame",
tooltip="First frame image from the original generation, if one was used.",
optional=True,
),
IO.Image.Input(
"last_frame",
tooltip="Last frame image from the original generation, if one was used.",
optional=True,
),
IO.Boolean.Input(
"watermark",
default=False,
tooltip="Whether to add an AIGC watermark to the video.",
advanced=True,
),
],
outputs=[
IO.Video.Output(),
],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type": "usd", "usd": 0.0715, "format": {"suffix": "/second"}}""",
),
)
@classmethod
async def execute(
cls,
model: dict,
video: Input.Video,
watermark: bool,
first_frame: torch.Tensor | None = None,
last_frame: torch.Tensor | None = None,
) -> IO.NodeOutput:
validate_string(model["prompt"], strip_whitespace=True, min_length=1)
try:
fps = float(video.get_frame_rate())
except Exception:
fps = 0.0
if fps and not (23.9 <= fps <= 24.1):
raise ValueError(
f"The source video is {fps:.2f} FPS. Regeneration accepts unmodified MiniMax H3 768P "
"outputs, which are 24 FPS."
)
try:
width, height = video.get_dimensions()
except Exception:
width = height = 0
if width and height and (width % 32 or height % 32 or width * height > 1_032_192):
raise ValueError(
f"The source video is {width}x{height}. Regeneration accepts MiniMax H3 768P outputs "
"(width and height divisible by 32, at most 1,032,192 total pixels); 2K outputs cannot "
"be used as a source."
)
try:
frame_count = video.get_frame_count()
except Exception:
frame_count = 0
if frame_count and (frame_count < 107 or frame_count > 362 or (frame_count - 107) % 17):
raise ValueError(
f"The source video has {frame_count} frames. Regeneration accepts unmodified "
"MiniMax H3 outputs, whose length is 107 to 362 frames in steps of 17 "
"(4 to 15 seconds at 24 FPS)."
)
reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
if (first_frame is not None or last_frame is not None) and (
reference_images or reference_videos or reference_audios
):
raise ValueError(
"First/last frame and reference media are mutually exclusive. Use frames for an "
"image-to-video prompt, or reference media for a reference-to-video prompt."
)
if reference_audios and not reference_images and not reference_videos:
raise ValueError("Reference audio cannot be used without a reference image or video.")
for frame in (first_frame, last_frame):
if frame is not None:
validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
validate_image_dimensions(frame, min_width=256, min_height=256)
for image in reference_images.values():
validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
validate_image_dimensions(image, min_width=256, min_height=256)
total_video_duration = 0.0
for i, ref_video in enumerate(reference_videos.values(), 1):
try:
ref_fps = float(ref_video.get_frame_rate())
except Exception:
ref_fps = 0.0
if ref_fps and not (23.9 <= ref_fps <= 60.5):
raise ValueError(f"Reference video {i} is {ref_fps:.2f} FPS. Supported range is 23.976-60 FPS.")
try:
dur = ref_video.get_duration()
except Exception:
continue
if dur < 1.8:
raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
total_video_duration += dur
if total_video_duration > 15.1:
raise ValueError(
f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
)
total_audio_duration = 0.0
for i, audio in enumerate(reference_audios.values(), 1):
dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
if dur < 1.8:
raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
total_audio_duration += dur
if total_audio_duration > 15.1:
raise ValueError(
f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
)
content: list = [
Hailuo03VideoContent(
video_url=Hailuo03VideoContentUrl(
url=await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video"),
),
role="base_video",
),
Hailuo03TextContent(text=model["prompt"]),
]
if first_frame is not None:
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, first_frame, max_images=1, wait_label="Uploading first frame"
)
)[0],
),
role="first_frame",
)
)
if last_frame is not None:
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, last_frame, max_images=1, wait_label="Uploading last frame"
)
)[0],
),
role="last_frame",
)
)
for i, image in enumerate(reference_images.values(), 1):
content.append(
Hailuo03ImageContent(
image_url=Hailuo03ImageContentUrl(
url=(
await upload_images_to_comfyapi(
cls, image, max_images=1, wait_label=f"Uploading image {i}"
)
)[0],
),
role="reference_image",
)
)
for i, ref_video in enumerate(reference_videos.values(), 1):
content.append(
Hailuo03VideoContent(
video_url=Hailuo03VideoContentUrl(
url=await upload_video_to_comfyapi(cls, ref_video, wait_label=f"Uploading video {i}"),
),
)
)
for audio in reference_audios.values():
content.append(
Hailuo03AudioContent(
audio_url=Hailuo03AudioContentUrl(
url=await upload_audio_to_comfyapi(
cls,
audio,
container_format="mp3",
codec_name="libmp3lame",
mime_type="audio/mpeg",
),
),
)
)
response = await sync_op(
cls,
ApiEndpoint(path=HAILUO_03_REGENERATION_ENDPOINT, method="POST"),
response_model=Hailuo03TaskCreationResponse,
data=Hailuo03RegenerationRequest(
model=HAILUO_03_MODELS[model["model"]],
content=content,
resolution=model["resolution"],
aigc_watermark=watermark,
),
)
task_result = await poll_op(
cls,
ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
response_model=Hailuo03TaskQueryResponse,
status_extractor=lambda r: r.task.status,
failed_statuses=HAILUO_03_FAILED_STATUSES,
poll_interval=10,
)
video_url = task_result.task.content.url if task_result.task.content else None
if not video_url:
raise Exception(f"No video URL in the response: {task_result.model_dump()}")
return IO.NodeOutput(await download_url_to_video_output(video_url))
class MinimaxExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@ -1541,8 +950,6 @@ class MinimaxExtension(ComfyExtension):
MinimaxHailuo03TextToVideoNode,
MinimaxHailuo03FirstLastFrameNode,
MinimaxHailuo03ReferenceNode,
MinimaxHailuo03ContextIRNode,
MinimaxHailuo03RegenerateNode,
]

View File

@ -27,7 +27,6 @@ from comfy_api_nodes.util import (
ApiEndpoint,
bytesio_to_image_tensor,
download_url_as_bytesio,
pad_images_to_common_channels,
resize_mask_to_image,
sync_op,
tensor_to_bytesio,
@ -622,7 +621,7 @@ class RecraftImageToImageNode(IO.ComfyNode):
images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
pbar.update(1)
return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
return IO.NodeOutput(torch.cat(images, dim=0))
class RecraftImageInpaintingNode(IO.ComfyNode):
@ -724,7 +723,7 @@ class RecraftImageInpaintingNode(IO.ComfyNode):
images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
pbar.update(1)
return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
return IO.NodeOutput(torch.cat(images, dim=0))
class RecraftTextToVectorNode(IO.ComfyNode):
@ -955,7 +954,7 @@ class RecraftReplaceBackgroundNode(IO.ComfyNode):
images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
pbar.update(1)
return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
return IO.NodeOutput(torch.cat(images, dim=0))
class RecraftRemoveBackgroundNode(IO.ComfyNode):
@ -996,7 +995,7 @@ class RecraftRemoveBackgroundNode(IO.ComfyNode):
image=image[i],
path="/proxy/recraft/images/removeBackground",
)
images.append(torch.cat([bytesio_to_image_tensor(x, mode="RGBA") for x in sub_bytes], dim=0))
images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
pbar.update(1)
images_tensor = torch.cat(images, dim=0)
@ -1048,7 +1047,7 @@ class RecraftCrispUpscaleNode(IO.ComfyNode):
images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
pbar.update(1)
return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
return IO.NodeOutput(torch.cat(images, dim=0))
class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode):

View File

@ -18,7 +18,6 @@ from .conversions import (
downscale_image_tensor_by_max_side,
downscale_video_to_max_pixels,
image_tensor_pair_to_batch,
pad_images_to_common_channels,
pil_to_bytesio,
resize_mask_to_image,
tensor_to_base64_string,
@ -93,7 +92,6 @@ __all__ = [
"downscale_image_tensor_by_max_side",
"downscale_video_to_max_pixels",
"image_tensor_pair_to_batch",
"pad_images_to_common_channels",
"pil_to_bytesio",
"resize_mask_to_image",
"tensor_to_base64_string",

View File

@ -16,14 +16,12 @@ from comfy_api.latest import Input, InputImpl, Types
from ._helpers import mimetype_to_extension
def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str | None = None) -> torch.Tensor:
def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor:
"""Converts image data from BytesIO to a torch.Tensor.
Args:
image_bytesio: BytesIO object containing the image data.
mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). Defaults
to RGBA when the decoded image carries transparency and RGB when it
does not, so an API that returns no alpha does not get an opaque one.
mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA").
Returns:
A torch.Tensor representing the image (1, H, W, C).
@ -33,8 +31,6 @@ def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str | None = None) ->
ValueError: If the specified mode is invalid.
"""
image = Image.open(image_bytesio)
if mode is None:
mode = "RGBA" if "A" in image.getbands() or "transparency" in image.info else "RGB"
image = image.convert(mode)
image_array = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(image_array).unsqueeze(0)
@ -57,17 +53,6 @@ def image_tensor_pair_to_batch(image1: torch.Tensor, image2: torch.Tensor) -> to
return torch.cat((image1, image2), dim=0)
def pad_images_to_common_channels(images: list[torch.Tensor]) -> list[torch.Tensor]:
"""Pads [B, H, W, C] image tensors with opaque alpha so they all share the largest channel count."""
channels = max(image.shape[-1] for image in images)
return [
torch.nn.functional.pad(image, (0, channels - image.shape[-1]), value=1.0)
if image.shape[-1] < channels
else image
for image in images
]
def tensor_to_bytesio(
image: torch.Tensor,
*,

View File

@ -718,7 +718,15 @@ class Noise_EmptyNoise:
self.seed = 0
def generate_noise(self, input_latent):
return comfy.sample.prepare_empty_noise(input_latent["samples"])
latent_image = input_latent["samples"]
if latent_image.is_nested:
tensors = latent_image.unbind()
zeros = []
for t in tensors:
zeros.append(torch.zeros(t.shape, dtype=t.dtype, layout=t.layout, device="cpu"))
return comfy.nested_tensor.NestedTensor(zeros)
else:
return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
class Noise_RandomNoise:

View File

@ -20,7 +20,6 @@ import comfy.model_sampling
import comfy.nested_tensor
import comfy.utils
import node_helpers
from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE
from comfy_api.latest import ComfyExtension, io
CANVAS_MULTIPLE = 32
@ -68,16 +67,6 @@ def _resize(image, width, height, crop):
return samples.movedim(1, -1)
def _encode_ref_audio(audio_vae, audio):
waveform = audio["waveform"] # [B, C, L]
sr = audio["sample_rate"]
vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
if sr != vae_sr:
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
z = audio_vae.encode(waveform[:1].movedim(1, -1)) # [1, 32, 2, T]
return z, z.shape[-1]
def _empty_av_latent(width, height, length, batch_size=1):
frame_count, latent_t, audio_t = temporal_shape(length)
video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
@ -155,87 +144,13 @@ class MiniMaxH3ImageToVideo(io.ComfyNode):
if keyframes:
for kf in keyframes:
kf["latent"] = vae.encode(kf.pop("image"))
cond = node_helpers.conditioning_set_values(cond, {"minimax_keyframes": keyframes})
cond = node_helpers.conditioning_set_values(cond, {
"minimax_keyframes": keyframes,
"minimax_frame_count": frame_count,
})
return io.NodeOutput(cond, latent)
class MiniMaxH3AddGuide(io.ComfyNode):
"""Anchor image and/or audio guides at an arbitrary pixel frame of the target video."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxH3AddGuide",
display_name="Add Guide for MiniMax H3",
category="model/conditioning/minimax",
description="Anchor an image, a short clip, audio, or a clip with its soundtrack at any frame of a MiniMax H3 video. Chain several nodes to anchor several frames.",
inputs=[
io.Conditioning.Input("positive"),
io.Vae.Input("vae", optional=True, tooltip="Video VAE, needed when an image is connected."),
io.Vae.Input("audio_vae", optional=True, tooltip="Audio VAE, needed when an audio is connected."),
io.Latent.Input("latent"),
io.Image.Input("image", optional=True, tooltip="Image or video frames to anchor. Multi-frame batches are anchored as a clip and cropped down to the model's valid clip lengths: 5, 22, 39... (17k + 5) frames. Batches shorter than 5 frames use only the first image."),
io.Audio.Input("audio", optional=True,
tooltip="Soundtrack to anchor starting at the same frame index, cropped to the video's remaining duration."),
io.Int.Input("frame_idx", default=0, min=-9999, max=9999,
tooltip="Frame index to anchor the image or the clip's first frame at. Negative values are counted from the end of the video."),
],
outputs=[io.Conditioning.Output(display_name="positive")],
)
@classmethod
def execute(cls, positive, latent, frame_idx, vae=None, audio_vae=None, image=None, audio=None) -> io.NodeOutput:
samples = latent["samples"]
if not samples.is_nested or len(samples.tensors) != 2 or samples.tensors[0].ndim != 5 or samples.tensors[0].shape[1] != 24:
raise ValueError("MiniMaxH3AddGuide expects a MiniMax H3 AV latent")
if image is None and audio is None:
raise ValueError("MiniMaxH3AddGuide needs an image or an audio to anchor")
video = samples.tensors[0]
height = video.shape[3] * 16
width = video.shape[4] * 16
frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(video.shape[2]))
guide_frames = 1
if image is not None:
if vae is None:
raise ValueError("anchoring guide frames needs the vae input")
guide_frames = image.shape[0]
if guide_frames < 5:
guide_frames = 1
else:
while guide_frames % 17 != 5:
guide_frames -= 1
resolved_frame_index = frame_idx if frame_idx >= 0 else frame_count + frame_idx
if resolved_frame_index < 0 or resolved_frame_index + guide_frames > frame_count:
if guide_frames == 1:
raise ValueError("frame_idx {} is outside the video's {} frames".format(frame_idx, frame_count))
raise ValueError("a {} frame guide clip at frame_idx {} does not fit in the video's {} frames".format(
guide_frames, frame_idx, frame_count))
keyframe = {"resolved_frame_index": resolved_frame_index}
if image is not None:
frames = _resize(image[:guide_frames], width, height, "center")
keyframe["latent"] = vae.encode(frames)
if audio is not None:
if audio_vae is None:
raise ValueError("anchoring guide audio needs the audio_vae input")
audio_latent, audio_rt = _encode_ref_audio(audio_vae, audio)
# the streams share one time axis: FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
max_rt = math.floor(samples.tensors[1].shape[-1] - FRAME_RESCALE * resolved_frame_index)
if max_rt < 1:
raise ValueError("frame_idx {} is past the end of the video's audio track".format(frame_idx))
if audio_rt > max_rt:
audio_latent = audio_latent[..., :max_rt].clone()
keyframe["audio_latent"] = audio_latent
keyframes = list(positive[0][1].get("minimax_keyframes", []))
keyframes.append(keyframe)
positive = node_helpers.conditioning_set_values(positive, {"minimax_keyframes": keyframes})
return io.NodeOutput(positive)
class MiniMaxH3ReferenceToVideo(io.ComfyNode):
"""ref2va: prompt + reference images / videos / audio -> conditioning + AV latent.
@ -282,6 +197,16 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
)
@staticmethod
def _encode_ref_audio(audio_vae, audio):
waveform = audio["waveform"] # [B, C, L]
sr = audio["sample_rate"]
vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
if sr != vae_sr:
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
z = audio_vae.encode(waveform[:1].movedim(1, -1)) # [1, 32, 2, T]
return z, z.shape[-1]
@classmethod
def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_size="match",
ref_images=None, ref_videos=None, ref_video_audios=None, ref_audios=None) -> io.NodeOutput:
@ -329,7 +254,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
z = vae.encode(frames)
audio_latent, ref_audio_t = (None, 0)
if soundtrack is not None:
audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, soundtrack)
audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, soundtrack)
# the soundtrack gets its own <Audio j> label, emitted before <Video k>
ref_items.append({"type": "audio"})
# Qwen sees the video at 2 fps with timestamps
@ -344,7 +269,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
for audio in (ref_audios or {}).values():
if audio is None:
continue
audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, audio)
audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, audio)
ref_items.append({"type": "audio"})
ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
@ -404,7 +329,6 @@ class MiniMaxH3Extension(ComfyExtension):
return [
EmptyMiniMaxH3LatentAV,
MiniMaxH3ImageToVideo,
MiniMaxH3AddGuide,
MiniMaxH3ReferenceToVideo,
MiniMaxH3SigmaShift,
]

View File

@ -1,77 +0,0 @@
import torch
from typing_extensions import override
import comfy.model_management
from comfy.ldm.minimax_music.ar import AUDIO_FRAMES_PER_SECOND, CFG_SCALE, CFG_TOP_K, C0_VOCAB_SIZE, MAX_AUDIO_FRAMES
from comfy.ldm.minimax_music.dit import latent_length
from comfy_api.latest import ComfyExtension, io
class MiniMaxMusic3TextEncode(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxMusic3TextEncode",
display_name="MiniMax Music3 Text Encode",
category="model/conditioning/minimax music",
description="Uses a MiniMax Music3 CLIP model to generate the acoustic conditioning sequence.",
inputs=[
io.Clip.Input("clip"),
io.String.Input("caption", multiline=True, dynamic_prompts=True),
io.String.Input("lyrics", multiline=True, dynamic_prompts=True),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
io.Float.Input("max_duration", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04, tooltip="Maximum duration in seconds; the model can end the song earlier."),
io.Float.Input("cfg_scale", default=CFG_SCALE, min=0.0, max=100.0, step=0.1, round=0.01, advanced=True),
io.Int.Input("top_k", default=CFG_TOP_K, min=1, max=C0_VOCAB_SIZE, advanced=True),
],
outputs=[
io.Conditioning.Output(),
io.Float.Output(display_name="seconds"),
],
)
@classmethod
def execute(cls, clip, caption, lyrics, seed, max_duration, cfg_scale, top_k):
max_audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(max_duration * AUDIO_FRAMES_PER_SECOND)))
tokens = clip.tokenize(caption, lyrics=lyrics, seed=seed, max_audio_frames=max_audio_frames, cfg_scale=cfg_scale, top_k=top_k)
conditioning = clip.encode_from_tokens_scheduled(tokens)
for cond in conditioning:
hidden = cond[0]
cond[1]["conditioning_scale"] = torch.ones((hidden.shape[0], 1, 1), device=hidden.device, dtype=hidden.dtype)
return io.NodeOutput(conditioning, conditioning[0][0].shape[1] / AUDIO_FRAMES_PER_SECOND)
class EmptyMiniMaxMusic3LatentAudio(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmptyMiniMaxMusic3LatentAudio",
display_name="Empty MiniMax Music3 Latent Audio",
category="model/latent/minimax music",
description="Creates an empty MiniMax Music3 audio latent for the requested duration.",
inputs=[
io.Float.Input("seconds", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04),
io.Int.Input("batch_size", default=1, min=1, max=4096),
],
outputs=[io.Latent.Output()],
)
@classmethod
def execute(cls, seconds, batch_size):
audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(seconds * AUDIO_FRAMES_PER_SECOND)))
latent = torch.zeros(
(batch_size, 128, latent_length(audio_frames)),
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
)
return io.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 512})
class MiniMaxMusic3Extension(ComfyExtension):
@override
async def get_node_list(self):
return [MiniMaxMusic3TextEncode, EmptyMiniMaxMusic3LatentAudio]
async def comfy_entrypoint():
return MiniMaxMusic3Extension()

View File

@ -29,7 +29,7 @@ class PreviewAny():
value = str(source)
elif source is not None:
try:
value = json.dumps(source, indent=4, ensure_ascii=False)
value = json.dumps(source, indent=4)
except Exception:
try:
value = str(source)

View File

@ -256,10 +256,12 @@ class TextGenerateLTX2Prompt(TextGenerate):
out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
# Drop reasoning, including a block left unclosed by max_length. Both system prompts ask
# for the original prompt back when there is nothing to give; empty conditions on nothing.
text = re.sub(r"<think>.*?(?:</think>|$)", "", out.args[0], flags=re.DOTALL).strip()
return io.NodeOutput(text or prompt)
text = out.args[0]
text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
if "</think>" in text: # unclosed/truncated reasoning: keep what follows the last close
text = text.rsplit("</think>", 1)[-1]
text = re.sub(r"</?think>|<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?", "", text).strip()
return io.NodeOutput(text)
class TextgenExtension(ComfyExtension):

View File

@ -1,3 +1,3 @@
# This file is automatically generated by the build process when version is
# updated in pyproject.toml.
__version__ = "0.33.0"
__version__ = "0.32.0"

12
main.py
View File

@ -248,17 +248,7 @@ import hook_breaker_ac10a0
import comfy.memory_management
import comfy.model_patcher
def dynamic_vram_supported():
if comfy.model_management.is_nvidia():
return True
if comfy.model_management.is_amd():
if comfy.model_management.rocm_version >= (7, 14):
return True
return False
if args.enable_dynamic_vram or (enables_dynamic_vram() and dynamic_vram_supported()):
if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()):
if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
else:

View File

@ -290,9 +290,6 @@ class ConditioningZeroOut:
conditioning_lyrics = d.get("conditioning_lyrics", None)
if conditioning_lyrics is not None:
d["conditioning_lyrics"] = torch.zeros_like(conditioning_lyrics)
conditioning_scale = d.get("conditioning_scale", None)
if conditioning_scale is not None:
d["conditioning_scale"] = torch.zeros_like(conditioning_scale)
n = [torch.zeros_like(t[0]), d]
c.append(n)
return (c, )
@ -1018,7 +1015,7 @@ class CLIPLoader:
CATEGORY = "model/loaders"
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm\nminimax: MiniMax H3 Qwen3-VL or Music3 Qwen/RVQ"
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
@ -1573,7 +1570,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
if disable_noise:
noise = comfy.sample.prepare_empty_noise(latent_image)
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
else:
batch_inds = latent["batch_index"] if "batch_index" in latent else None
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
@ -2452,7 +2449,6 @@ async def init_builtin_extra_nodes():
"nodes_mahiro.py",
"nodes_lt_upsampler.py",
"nodes_lt_audio.py",
"nodes_minimax_music.py",
"nodes_minimax_h3.py",
"nodes_lt.py",
"nodes_hooks.py",

View File

@ -35,10 +35,6 @@ components:
description: Timestamp when the asset was last accessed
format: date-time
type: string
loader_path:
description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
nullable: true
type: string
metadata:
additionalProperties: true
description: System-managed metadata from download sources (HuggingFace, CivitAI, etc.) - read-only, not user-modifiable
@ -169,10 +165,6 @@ components:
format: uuid
nullable: true
type: string
loader_path:
description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
nullable: true
type: string
mime_type:
description: Updated MIME type of the asset
type: string
@ -196,31 +188,6 @@ components:
- id
- updated_at
type: object
ChurnkeyAuthResponse:
description: |
Credentials the Churnkey embed requires to launch the cancel flow.
`auth_hash` is hex-encoded HMAC-SHA256 of `customer_id` signed with the
server's CHURNKEY_HMAC_SECRET; it is bound to that single customer ID
and must not be reused for other customers.
properties:
auth_hash:
description: Hex-encoded HMAC-SHA256(customer_id, CHURNKEY_HMAC_SECRET)
type: string
customer_id:
description: Stripe customer ID for the workspace
type: string
mode:
description: Churnkey environment matching the configured app
enum:
- live
- test
- sandbox
type: string
required:
- customer_id
- auth_hash
- mode
type: object
CreateWorkflowRequest:
description: Request body for creating a new saved workflow.
properties:
@ -544,25 +511,6 @@ components:
required:
- history
type: object
JobAssetsResponse:
description: Paginated list of the assets produced by a single job.
properties:
assets:
description: The job's output assets for the requested page (empty when the job produced none)
items:
$ref: '#/components/schemas/JobOutputAsset'
type: array
job_id:
description: ID of the job these assets belong to
format: uuid
type: string
pagination:
$ref: '#/components/schemas/PaginationInfo'
required:
- job_id
- assets
- pagination
type: object
JobCancelResponse:
description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
properties:
@ -617,9 +565,6 @@ components:
additionalProperties: true
description: Primary preview output (only for terminal states)
type: object
previewable_outputs_count:
description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
type: integer
status:
description: User-friendly job status
enum:
@ -652,13 +597,6 @@ components:
workflow_id:
description: UUID identifying the workflow graph definition
type: string
workflow_version_id:
description: |
UUID of the cloud workflow version this job is pinned to, if the
submission carried one (see PromptRequest's workflow_version_id).
Absent for jobs submitted without that association, including
every job submitted through the public API v2 today.
type: string
workspace_id:
description: |
ID of the workspace that owns this job. A successful (200)
@ -707,9 +645,6 @@ components:
additionalProperties: true
description: Primary preview output (only present for terminal states)
type: object
previewable_outputs_count:
description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
type: integer
status:
description: User-friendly job status
enum:
@ -727,56 +662,6 @@ components:
- status
- create_time
type: object
JobOutputAsset:
description: |
An asset produced by a job, enriched with the per-output node context
(`node_id`, `output_key`, `output_index`) correlated from the job's
execution outputs by content hash. The node-context fields are null
when the asset cannot be matched to an output entry.
properties:
created_at:
description: Timestamp when the asset was created
format: date-time
type: string
hash:
description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$
type: string
id:
description: Unique identifier for the asset
format: uuid
type: string
mime_type:
description: MIME type of the asset
type: string
name:
description: Name of the asset file
type: string
node_id:
description: ID of the workflow node that produced this asset, if known
nullable: true
type: string
output_index:
description: Zero-based index of this asset within the node's output slot, if known
nullable: true
type: integer
output_key:
description: Output slot key under the producing node (e.g. "images"), if known
nullable: true
type: string
preview_url:
description: Relative URL for asset preview/thumbnail
format: uri-reference
type: string
size:
description: Size of the asset in bytes
format: int64
type: integer
required:
- id
- name
- created_at
type: object
JobStatusResponse:
description: Job status information
properties:
@ -1637,11 +1522,7 @@ paths:
operationId: listAssets
parameters:
- deprecated: true
description: |
Deprecated alias for `tags_all`, kept permanently for existing
callers. Filter assets that have ALL of these tags. Combining it
with `tags_all`, or exceeding 100 tags (counted after removing
empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
explode: false
in: query
name: include_tags
@ -1651,11 +1532,7 @@ paths:
type: array
style: form
- deprecated: true
description: |
Deprecated alias for `tags_none`, kept permanently for existing
callers. Exclude assets that have ANY of these tags. Combining it
with `tags_none`, or exceeding 100 tags (counted after removing
empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
explode: false
in: query
name: exclude_tags
@ -1664,14 +1541,7 @@ paths:
type: string
type: array
style: form
- description: |
Filter assets that have ALL of these tags. Tag values are opaque
byte-strings compared exactly and case-sensitively; unknown tags
are not an error — they simply match nothing. Replaces the
deprecated `include_tags`. Sending both spellings, listing the
same tag here and in `tags_none`, or exceeding 100 tags per list
(counted after removing empty values and duplicates) returns 400
`INVALID_TAG_FILTER`.
- description: Filter assets that have ALL of these tags
explode: false
in: query
name: tags_all
@ -1680,14 +1550,7 @@ paths:
type: string
type: array
style: form
- description: |
Filter assets that have AT LEAST ONE of these tags. Combines with
`tags_all`/`tags_none` by intersection (`tags_none` always wins;
overlap with `tags_none` is allowed and leaves a dead term).
Supplying a positive tag filter (`tags_any`, `tags_all`, or
`include_tags`) replaces the default category filter that is
otherwise applied. Lists over 100 tags (counted after removing
empty values and duplicates) return 400 `INVALID_TAG_FILTER`.
- description: Filter assets that have AT LEAST ONE of these tags
explode: false
in: query
name: tags_any
@ -1696,11 +1559,7 @@ paths:
type: string
type: array
style: form
- description: |
Exclude assets that have ANY of these tags. Replaces the
deprecated `exclude_tags`. Sending both spellings, or exceeding
100 tags per list (counted after removing empty values and
duplicates), returns 400 `INVALID_TAG_FILTER`.
- description: Exclude assets that have ANY of these tags
explode: false
in: query
name: tags_none
@ -2483,11 +2342,7 @@ paths:
operationId: getAssetTagHistogram
parameters:
- deprecated: true
description: |
Deprecated alias for `tags_all`, kept permanently for existing
callers. Filter assets that have ALL of these tags. The same
combination and list-size rules as on `/api/assets` apply
(400 `INVALID_TAG_FILTER`).
description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
explode: false
in: query
name: include_tags
@ -2497,11 +2352,7 @@ paths:
type: array
style: form
- deprecated: true
description: |
Deprecated alias for `tags_none`, kept permanently for existing
callers. Exclude assets that have ANY of these tags. The same
combination and list-size rules as on `/api/assets` apply
(400 `INVALID_TAG_FILTER`).
description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
explode: false
in: query
name: exclude_tags
@ -2510,10 +2361,7 @@ paths:
type: string
type: array
style: form
- description: |
Filter assets that have ALL of these tags. Replaces the deprecated
`include_tags`. The same combination and list-size rules as on
`/api/assets` apply (400 `INVALID_TAG_FILTER`).
- description: Filter assets that have ALL of these tags
explode: false
in: query
name: tags_all
@ -2522,11 +2370,7 @@ paths:
type: string
type: array
style: form
- description: |
Filter assets that have AT LEAST ONE of these tags. Combines with
`tags_all`/`tags_none` by intersection (`tags_none` always wins).
The same combination and list-size rules as on `/api/assets` apply
(400 `INVALID_TAG_FILTER`).
- description: Filter assets that have AT LEAST ONE of these tags
explode: false
in: query
name: tags_any
@ -2535,10 +2379,7 @@ paths:
type: string
type: array
style: form
- description: |
Exclude assets that have ANY of these tags. Replaces the deprecated
`exclude_tags`. The same combination and list-size rules as on
`/api/assets` apply (400 `INVALID_TAG_FILTER`).
- description: Exclude assets that have ANY of these tags
explode: false
in: query
name: tags_none
@ -2599,49 +2440,6 @@ paths:
summary: Get tag histogram for filtered assets
tags:
- file
/api/billing/churnkey/auth:
get:
description: |
Returns the Stripe customer identifier and a server-signed
HMAC-SHA256 of the customer ID, used to launch the Churnkey-hosted
cancellation flow embed.
operationId: getChurnkeyAuth
responses:
"200":
content:
application/json:
schema:
$ref: '#/components/schemas/ChurnkeyAuthResponse'
description: Success
"401":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Unauthorized
"404":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Workspace has no Stripe customer (never subscribed)
"500":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Internal server error
"503":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Churnkey is not configured on the server
security:
- BearerAuth: []
summary: Get Churnkey HMAC auth credentials
tags:
- billing
/api/embeddings:
get:
description: Returns the list of text-encoder embeddings available on disk.
@ -2662,10 +2460,9 @@ paths:
Returns a list of model folders available in the system.
This is an experimental endpoint that replaces the legacy /models endpoint.
Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
The folder vocabulary is resolved per request from the caller's identity: where the
model_type migration is active for that caller the names are model_type folder_names
(e.g. `ultralytics_bbox`), and a folder with no folder_name mapping is returned by its
directory path. An authenticated response can therefore differ from an anonymous one.
Once the model_type migration is active the names are model_type folder_names
(e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
its directory path.
operationId: getModelFolders
responses:
"200":
@ -2682,10 +2479,7 @@ paths:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Internal server error
security:
- ApiKeyAuth: []
- BearerAuth: []
- {}
security: []
summary: Get available model folders
tags:
- file
@ -2694,10 +2488,6 @@ paths:
description: |
Returns a list of models available in the specified folder.
This is an experimental endpoint that provides enhanced model information.
Accepted folder identifiers are those returned by /api/experiment/models for the same
caller. That vocabulary is request-scoped, so list folders and fetch a folder's models
with the same credentials — a name obtained anonymously may not resolve when
authenticated, and vice versa.
operationId: getModelsInFolder
parameters:
- description: The folder name to list models from
@ -2728,10 +2518,7 @@ paths:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Internal server error
security:
- ApiKeyAuth: []
- BearerAuth: []
- {}
security: []
summary: Get models in a specific folder
tags:
- file
@ -3368,74 +3155,6 @@ paths:
summary: Get full job details
tags:
- workflow
/api/jobs/{job_id}/assets:
get:
description: |
Retrieve a paginated list of the assets produced by a specific job,
enriched with the per-output node context (`node_id`, `output_key`,
`output_index`) correlated from the job's execution outputs by content
hash. Unlike `GET /api/assets?job_ids={id}`, this endpoint is scoped to a
single job and carries node-level placement, making it suited to job
output views rather than the general asset browser. Returns an empty
`assets` array for jobs that produced no assets.
operationId: getJobAssets
parameters:
- description: Job identifier (UUID)
in: path
name: job_id
required: true
schema:
format: uuid
type: string
- description: Maximum number of assets to return (1-500)
in: query
name: limit
schema:
default: 20
maximum: 500
minimum: 1
type: integer
- description: Number of assets to skip for pagination
in: query
name: offset
schema:
default: 0
minimum: 0
type: integer
responses:
"200":
content:
application/json:
schema:
$ref: '#/components/schemas/JobAssetsResponse'
description: Success - Job assets returned
"400":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Invalid request parameters
"401":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Unauthorized - Authentication required
"404":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Job not found or does not belong to the user
"500":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Internal server error
summary: List a job's output assets
tags:
- workflow
/api/jobs/{job_id}/cancel:
post:
description: |
@ -3641,12 +3360,6 @@ paths:
schema:
$ref: '#/components/schemas/PromptErrorResponse'
description: Payment required - Insufficient credits
"403":
content:
application/json:
schema:
$ref: '#/components/schemas/PromptErrorResponse'
description: Workspace governance policy blocks one or more partner providers (error.type PARTNER_NODE_DISABLED; error.class_types lists the offending nodes, error.providers the disabled providers)
"413":
content:
application/json:
@ -3658,7 +3371,7 @@ paths:
application/json:
schema:
$ref: '#/components/schemas/PromptErrorResponse'
description: 'Retryable backpressure. Two distinct causes, disambiguated by the body''s `error.type`, NOT by parsing `error.message`: `PAYMENT_REQUIRED` / `FREE_TIER_UNAVAILABLE` / `FREE_TIER_EXHAUSTED` / `PARTNER_NODE_PAYMENT_REQUIRED` (a billing gate - retrying without paying never succeeds), or `QUEUE_LIMIT` (this workspace''s bounded job queue is full - retrying after some queued jobs complete will succeed).'
description: Payment required - User has not paid
"500":
content:
application/json:
@ -5497,8 +5210,6 @@ tags:
name: user
- description: Background task management
name: task
- description: Workspace billing and subscription management
name: billing
- description: Workflow storage and version management
name: workflows
- description: Job queue state and control

View File

@ -1,6 +1,6 @@
[project]
name = "ComfyUI"
version = "0.33.0"
version = "0.32.0"
readme = "README.md"
license = { file = "LICENSE" }
requires-python = ">=3.10"

View File

@ -1,6 +1,6 @@
comfyui-frontend-package==1.49.6
comfyui-workflow-templates==0.11.41
comfyui-embedded-docs==0.5.10
comfyui-frontend-package==1.48.7
comfyui-workflow-templates==0.11.39
comfyui-embedded-docs==0.5.9
torch
torchsde
torchvision
@ -22,7 +22,7 @@ alembic
SQLAlchemy>=2.0.0
filelock
av>=16.0.0
comfy-kitchen==0.2.31
comfy-kitchen==0.2.30
comfy-aimdo==0.4.13
requests
simpleeval>=1.0.0

View File

@ -1,57 +0,0 @@
import asyncio
import base64
from io import BytesIO
import torch
from PIL import Image
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
from comfy_api_nodes.apis.gemini import ( # noqa: E402
GeminiCandidate,
GeminiContent,
GeminiGenerateContentResponse,
GeminiInlineData,
GeminiPart,
)
from comfy_api_nodes.nodes_gemini import get_image_from_response # noqa: E402
def image_part(mode, color):
buffer = BytesIO()
Image.new(mode, (4, 4), color).save(buffer, format="PNG")
return GeminiPart(
inlineData=GeminiInlineData(
data=base64.b64encode(buffer.getvalue()).decode(),
mimeType="image/png",
)
)
def response(*parts):
return GeminiGenerateContentResponse(
candidates=[GeminiCandidate(content=GeminiContent(parts=list(parts), role="model"))]
)
def test_rgb_only_response_stays_three_channels():
out = asyncio.run(get_image_from_response(response(image_part("RGB", (10, 20, 30)))))
assert out.shape == (1, 4, 4, 3)
def test_mixed_rgb_and_rgba_parts_are_padded_to_the_same_width():
out = asyncio.run(
get_image_from_response(
response(
image_part("RGB", (10, 20, 30)),
image_part("RGBA", (10, 20, 30, 0)),
)
)
)
assert out.shape == (2, 4, 4, 4)
# the part that had no alpha is padded opaque, the transparent one is preserved
assert out[0, ..., 3].min() == 1.0
assert out[1, ..., 3].max() == 0.0

View File

@ -1,80 +0,0 @@
from io import BytesIO
import pytest
import torch
from PIL import Image
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
from comfy_api_nodes.util.conversions import bytesio_to_image_tensor, pad_images_to_common_channels # noqa: E402
def encode(image: Image.Image, image_format: str = "PNG") -> BytesIO:
buffer = BytesIO()
image.save(buffer, format=image_format)
buffer.seek(0)
return buffer
def test_rgb_png_stays_three_channels():
tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30))))
assert tensor.shape == (1, 4, 4, 3)
def test_jpeg_stays_three_channels():
tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30)), "JPEG"))
assert tensor.shape == (1, 4, 4, 3)
def test_grayscale_is_expanded_to_rgb():
tensor = bytesio_to_image_tensor(encode(Image.new("L", (4, 4), 128)))
assert tensor.shape == (1, 4, 4, 3)
def test_rgba_png_keeps_its_alpha():
tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 0))))
assert tensor.shape == (1, 4, 4, 4)
assert tensor[..., 3].max() == 0.0
def test_palette_png_with_transparency_keeps_its_alpha():
image = Image.new("P", (4, 4), 1)
image.putpalette([0, 0, 0, 255, 255, 255])
image.info["transparency"] = 0
image.putpixel((0, 0), 0)
tensor = bytesio_to_image_tensor(encode(image))
assert tensor.shape == (1, 4, 4, 4)
assert tensor[0, 0, 0, 3] == 0.0
assert tensor[0, 1, 1, 3] == 1.0
@pytest.mark.parametrize("mode,channels", [("RGB", 3), ("RGBA", 4)])
def test_explicit_mode_is_respected(mode, channels):
tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 128))), mode=mode)
assert tensor.shape == (1, 4, 4, channels)
def test_pad_mixed_channels_concatenates():
rgb = torch.rand(1, 4, 4, 3)
rgba = torch.rand(2, 4, 4, 4)
padded = pad_images_to_common_channels([rgb, rgba])
result = torch.cat(padded, dim=0)
assert result.shape == (3, 4, 4, 4)
def test_pad_adds_opaque_alpha_and_keeps_rgb_values():
rgb = torch.rand(1, 4, 4, 3)
rgba = torch.rand(1, 4, 4, 4)
padded_rgb, padded_rgba = pad_images_to_common_channels([rgb, rgba])
assert torch.equal(padded_rgb[..., :3], rgb)
assert padded_rgb[..., 3].min() == 1.0
assert padded_rgba is rgba
def test_pad_leaves_homogeneous_channels_unchanged():
images = [torch.rand(1, 4, 4, 3), torch.rand(2, 4, 4, 3)]
padded = pad_images_to_common_channels(images)
assert all(p is i for p, i in zip(padded, images))

View File

@ -187,7 +187,7 @@ class TestMathExpressionExecute:
self._exec("a / b", a=1, b=0)
def test_sqrt_negative_raises(self):
with pytest.raises(ValueError, match="math domain error|expected a nonnegative input"):
with pytest.raises(ValueError, match="math domain error"):
self._exec("sqrt(a)", a=-1)
def test_overflow_inf_raises(self):

View File

@ -1,30 +0,0 @@
from unittest.mock import patch, MagicMock
mock_nodes = MagicMock()
mock_nodes.MAX_RESOLUTION = 16384
mock_server = MagicMock()
with patch.dict("sys.modules", {"nodes": mock_nodes, "server": mock_server}):
from comfy_extras.nodes_preview_any import PreviewAny
class TestPreviewAnyMain:
@staticmethod
def _exec(source) -> dict:
return PreviewAny().main(source)
def test_dict_keeps_non_ascii(self):
result = self._exec({"greeting": "你好"})
assert "你好" in result["ui"]["text"][0]
assert "\\u" not in result["ui"]["text"][0]
assert result["result"][0] == result["ui"]["text"][0]
def test_list_keeps_non_ascii(self):
result = self._exec(["你好", "こんにちは"])
assert "こんにちは" in result["result"][0]
assert "\\u" not in result["result"][0]
def test_string_passthrough(self):
result = self._exec("你好")
assert result["ui"]["text"][0] == "你好"
assert result["result"][0] == "你好"

View File

@ -1,61 +0,0 @@
"""Gemma4 chat template regression tests."""
import pytest
import torch
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
import comfy.text_encoders.gemma4 as gemma4 # noqa: E402
PROMPT = "describe a cute anime girl with fennec ears"
THOUGHT_BLOCK = "<|channel>thought\n<channel|>"
# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a
# closed thought block when thinking is off.
NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B]
PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B]
class _CaptureTemplate:
"""Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files."""
llama_text = ""
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
self.llama_text = text
return {}
def build_template(variant, **kwargs):
prime = variant.tokenizer.tokenizer_class.prime_empty_thought
probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})()
probe.tokenize_with_weights(PROMPT, **kwargs)
return probe.llama_text
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
def test_thinking_enabled_only_asks_via_the_system_turn(variant):
template = build_template(variant, skip_template=False, thinking=True)
assert template == f"<|turn>system\n<|think|>\n<turn|>\n<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
@pytest.mark.parametrize("variant", NO_PRIMING)
def test_thinking_disabled_does_not_prime_a_thought_channel(variant):
template = build_template(variant, skip_template=False, thinking=False)
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
assert "channel" not in template
assert "<|think|>" not in template
@pytest.mark.parametrize("variant", PRIMING)
def test_thinking_disabled_primes_a_thought_channel(variant):
template = build_template(variant, skip_template=False, thinking=False)
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n{THOUGHT_BLOCK}"
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
@pytest.mark.parametrize("thinking", [False, True])
def test_skip_template_passes_text_through_unchanged(variant, thinking):
assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT