diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 9de2440875f..c6660846d06 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -180,6 +180,7 @@ def from_string(cls, value: str): 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.") diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index c4270022beb..dc737fc7d27 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -957,6 +957,11 @@ 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 diff --git a/comfy/ldm/minimax_music/__init__.py b/comfy/ldm/minimax_music/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/comfy/ldm/minimax_music/ar.py b/comfy/ldm/minimax_music/ar.py new file mode 100644 index 00000000000..28215cb5832 --- /dev/null +++ b/comfy/ldm/minimax_music/ar.py @@ -0,0 +1,337 @@ +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, dtype, device, operations): + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.merged_qkv = None + self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device) + 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, dtype, device, operations): + super().__init__() + self.merged_mlp = None + self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device) + 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, dtype, device, operations): + super().__init__() + self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device) + self.self_attn = RVQAttention(hidden_size, num_heads, dtype, device, operations) + self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device) + self.mlp = RVQMLP(hidden_size, intermediate_size, 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"]) + 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"]), + 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 + qwen_config.merged_qkv = None + qwen_config.merged_mlp = None + 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") diff --git a/comfy/ldm/minimax_music/dav.py b/comfy/ldm/minimax_music/dav.py new file mode 100644 index 00000000000..d442559f4d4 --- /dev/null +++ b/comfy/ldm/minimax_music/dav.py @@ -0,0 +1,137 @@ +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 diff --git a/comfy/ldm/minimax_music/dit.py b/comfy/ldm/minimax_music/dit.py new file mode 100644 index 00000000000..211e0d7dbd9 --- /dev/null +++ b/comfy/ldm/minimax_music/dit.py @@ -0,0 +1,213 @@ +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 diff --git a/comfy/ldm/minimax_music/prompt.py b/comfy/ldm/minimax_music/prompt.py new file mode 100644 index 00000000000..5f197ee123f --- /dev/null +++ b/comfy/ldm/minimax_music/prompt.py @@ -0,0 +1,70 @@ +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"(?<|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}") diff --git a/comfy/model_base.py b/comfy/model_base.py index 7d855f5a1ba..90cab7ac0a7 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -22,6 +22,7 @@ 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 @@ -2337,6 +2338,18 @@ def extra_conds(self, **kwargs): 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) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index aec2052909b..e4bf30b7869 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -44,6 +44,13 @@ 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] diff --git a/comfy/model_management.py b/comfy/model_management.py index 15c03dc771d..ff963eb8ea8 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1368,9 +1368,14 @@ def current_stream(device): 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 @@ -1425,6 +1430,10 @@ 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(): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index cb44e7394e0..72942aa0448 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -1887,8 +1887,29 @@ def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False 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 = x + *_, module_mem, n, m, params, end_of_block = x def set_dirty(item, dirty): if dirty or not hasattr(item, "_v_signature"): @@ -1981,6 +2002,13 @@ def force_load_param(self, param_key, device_to): 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 diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index aa6d22d77eb..2aad5eea79f 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -1,11 +1,18 @@ +import torch +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(comfy_modules): +def cleanup_prefetched_modules(module, comfy_modules): for s in comfy_modules: prefetch = getattr(s, "_prefetch", None) if prefetch is None: @@ -17,39 +24,74 @@ def cleanup_prefetched_modules(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 cleanup_prefetch_queues(): - global PREFETCH_QUEUES + global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS for queue in PREFETCH_QUEUES: for entry in queue: if entry is None or not isinstance(entry, tuple): continue _, prefetch_state = entry - comfy_modules = prefetch_state[1] + prefetched_module, comfy_modules = prefetch_state if comfy_modules is not None: - cleanup_prefetched_modules(comfy_modules) + cleanup_prefetched_modules(prefetched_module, comfy_modules) PREFETCH_QUEUES = [] + for module in GRAPH_MODULES: + del module._comfy_graph + GRAPH_MODULES.clear() + GRAPH_WARMED_MODULES.clear() + GRAPH_CAPTURE_STREAMS = {} -def prefetch_queue_pop(queue, device, module): +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) 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)) - _, comfy_modules = prefetch_state + prefetched_module, comfy_modules = prefetch_state if comfy_modules is not None: - cleanup_prefetched_modules(comfy_modules) + cleanup_prefetched_modules(prefetched_module, 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 = [] - for s in prefetch.modules(): - if hasattr(s, "_v"): - comfy_modules.append(s) + 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) registerable_size = 0 for s in comfy_modules: @@ -59,11 +101,41 @@ def prefetch_queue_pop(queue, device, module): if lowvram_fn is not None: registerable_size += lowvram_fn.memory_required() - offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True) + offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=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) - queue[0] = (offload_stream, (prefetch, comfy_modules)) + 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: + 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) def make_prefetch_queue(queue, device, transformer_options): if (not transformer_options.get("prefetch_dynamic_vbars", False) diff --git a/comfy/ops.py b/comfy/ops.py index 9ec44cfa274..73ae4667475 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -123,10 +123,12 @@ 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): +def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False): 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 @@ -163,6 +165,8 @@ def get_cast_buffer(buffer_size): 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, @@ -255,10 +259,12 @@ def handle_pin(m, pin, source, dest, subset="weights", size=None): 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): +def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True): prefetch = getattr(s, "_prefetch", None) @@ -298,7 +304,7 @@ def to_dequant(tensor, dtype): tensor = tensor.dequantize() return tensor - if orig.dtype != dtype or len(fns) > 0: + if (return_weights and 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) @@ -325,7 +331,7 @@ def to_dequant(tensor, dtype): if prefetch["signature"] is not None: prefetch["resident"] = True - return weight, bias + return (weight, bias) if return_weights else None def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False): diff --git a/comfy/sd.py b/comfy/sd.py index 46c9acba12f..94f4f284f58 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -25,6 +25,7 @@ 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 @@ -32,6 +33,7 @@ import os import comfy.utils +import comfy.ops from . import clip_vision from . import gligen @@ -74,6 +76,7 @@ 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 @@ -515,7 +518,22 @@ def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None) self.audio_sample_rate = 44100 if config is None: - if "decoder.mid.block_1.mix_factor" in sd: + 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: 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] @@ -1692,7 +1710,15 @@ class EmptyClass: clip_target.params = {} if len(clip_data) == 1: te_model = detect_te_model(clip_data[0]) - if te_model == TEModel.CLIP_G: + 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.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel + clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer + elif te_model == TEModel.CLIP_G: if clip_type == CLIPType.STABLE_CASCADE: clip_target.clip = sdxl_clip.StableCascadeClipModel clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer diff --git a/comfy/supported_models.py b/comfy/supported_models.py index b9952db5509..d6d3c857f8c 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -16,6 +16,7 @@ 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 @@ -2200,6 +2201,25 @@ def clip_target(self, state_dict={}): 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={}): + return supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel) + class LongCatImage(supported_models_base.BASE): unet_config = { @@ -2494,6 +2514,7 @@ def get_model(self, state_dict, prefix="", device=None): ChromaRadiance, ACEStep, ACEStep15, + MiniMaxMusic3, Omnigen2, Boogu, MageFlow, diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index 371ec1bbc95..ddac547e050 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -5,15 +5,40 @@ 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 + +def detect_merged_config(state_dict, prefix="", layer_prefix="model.layers.0."): + return { + "merged_qkv": "{}{}self_attn.qkv_proj.weight".format(prefix, layer_prefix) in state_dict, + "merged_mlp": "{}{}mlp.gate_up_proj.weight".format(prefix, layer_prefix) in state_dict, + } + + +@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 @@ -249,6 +274,9 @@ 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 @@ -498,9 +526,14 @@ def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = Non self.inner_size = self.num_heads * self.head_dim ops = ops or nn - 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.kv_size = self.num_kv_heads * self.head_dim + self.merged_qkv = getattr(config, "merged_qkv", False) + if self.merged_qkv is not False: + self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype) + if self.merged_qkv is not True: + 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.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype) self.q_norm = None @@ -522,9 +555,12 @@ def forward( ): batch_size, seq_length, _ = hidden_states.shape - xq = self.q_proj(hidden_states) - xk = self.k_proj(hidden_states) - xv = self.v_proj(hidden_states) + 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 = 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) @@ -537,8 +573,29 @@ def forward( xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis) - present_key_value = None - if past_key_value is not None: + 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: index = 0 num_tokens = xk.shape[2] if len(past_key_value) > 0: @@ -569,15 +626,27 @@ 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.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.merged_mlp = getattr(config, "merged_mlp", False) + if self.merged_mlp is not False: + self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype) + if self.merged_mlp is not True: + 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): @@ -596,6 +665,7 @@ def forward( optimized_attention=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): + output = x # Self Attention residual = x x = self.input_layernorm(x) @@ -612,7 +682,7 @@ def forward( residual = x x = self.post_attention_layernorm(x) x = self.mlp(x) - x = residual + x + x = torch.add(residual, x, out=output) return x, present_key_value @@ -641,6 +711,7 @@ def forward( 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: @@ -676,7 +747,7 @@ def forward( x = self.pre_feedforward_layernorm(x) x = self.mlp(x) x = self.post_feedforward_layernorm(x) - x = residual + x + x = torch.add(residual, x, out=output) return x, present_key_value @@ -688,9 +759,14 @@ def forward(self, input_ids, out_dtype=None): 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": @@ -713,8 +789,27 @@ def __init__(self, config, device=None, dtype=None, ops=None): 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): - return past_key_values[0][2] + 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 def compute_freqs_cis(self, position_ids, device): return precompute_freqs_cis(self.config.head_dim, @@ -756,6 +851,33 @@ def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermed 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 @@ -769,7 +891,8 @@ def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermed elif intermediate_output < 0: intermediate_output = len(self.layers) + intermediate_output - next_key_values = [] + 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 [] for i, layer in enumerate(self.layers): if all_intermediate is not None: if only_layers is None or (i in only_layers): @@ -779,16 +902,24 @@ def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermed if past_key_values is not None: past_kv = past_key_values[i] if len(past_key_values) > 0 else [] - x, current_kv = layer( - x=x, - attention_mask=mask, - freqs_cis=freqs_cis, - optimized_attention=optimized_attention, - past_key_value=past_kv, - ) - - if current_kv is not None: - next_key_values.append(current_kv) + if fixed_kv: + past_kv.prepare(seq_len) + + 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) # 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): @@ -797,6 +928,9 @@ def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermed 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) @@ -810,7 +944,7 @@ def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermed if intermediate is not None and final_layer_norm_intermediate and self.norm is not None: intermediate = self.norm(intermediate) - if len(next_key_values) > 0: + if next_key_values: return x, intermediate, next_key_values else: return x, intermediate @@ -874,12 +1008,7 @@ def logits(self, x): return torch.nn.functional.linear(input, weight, None) def init_kv_cache(self, 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 + return self.model.init_kv_cache(batch, max_cache_len, device, execution_dtype) 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 diff --git a/comfy/text_encoders/minimax_music.py b/comfy/text_encoders/minimax_music.py new file mode 100644 index 00000000000..37d072664d1 --- /dev/null +++ b/comfy/text_encoders/minimax_music.py @@ -0,0 +1,129 @@ +import torch +from tokenizers import Tokenizer + +import comfy.ops +import comfy.text_encoders.llama +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, +} + + +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={}): + 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, 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): + def select_projections(layers, config): + for layer in layers: + if layer.self_attn.merged_qkv is None: + if config["merged_qkv"]: + del layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj + else: + del layer.self_attn.qkv_proj + layer.self_attn.merged_qkv = config["merged_qkv"] + if layer.mlp.merged_mlp is None: + if config["merged_mlp"]: + del layer.mlp.gate_proj, layer.mlp.up_proj + else: + del layer.mlp.gate_up_proj + layer.mlp.merged_mlp = config["merged_mlp"] + + select_projections(self.model.layers, comfy.text_encoders.llama.detect_merged_config(state_dict)) + select_projections( + self.model.audio_decoder.layers, + comfy.text_encoders.llama.detect_merged_config(state_dict, layer_prefix="model.audio_decoder.layers.0."), + ) + 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)) diff --git a/comfy_api_nodes/apis/bria.py b/comfy_api_nodes/apis/bria.py index 7a98428c3f3..f55de74bc24 100644 --- a/comfy_api_nodes/apis/bria.py +++ b/comfy_api_nodes/apis/bria.py @@ -57,6 +57,81 @@ 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(...) @@ -72,6 +147,26 @@ 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(...) diff --git a/comfy_api_nodes/apis/minimax.py b/comfy_api_nodes/apis/minimax.py index bac4572d417..12a0853ac02 100644 --- a/comfy_api_nodes/apis/minimax.py +++ b/comfy_api_nodes/apis/minimax.py @@ -161,12 +161,30 @@ class Hailuo03TaskCreationRequest(BaseModel): ..., min_length=1 ) resolution: str = Field(...) - duration: int = Field(..., ge=5, le=15) + duration: int = Field(..., ge=4, 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(...) @@ -178,6 +196,7 @@ class Hailuo03TaskError(BaseModel): class Hailuo03TaskContent(BaseModel): url: str | None = Field(None) + prompt: str | None = Field(None) class Hailuo03TaskUsage(BaseModel): diff --git a/comfy_api_nodes/nodes_bria.py b/comfy_api_nodes/nodes_bria.py index 77f780a3bbc..90cade2d06a 100644 --- a/comfy_api_nodes/nodes_bria.py +++ b/comfy_api_nodes/nodes_bria.py @@ -6,7 +6,13 @@ 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, @@ -21,13 +27,30 @@ 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): @@ -243,6 +266,503 @@ async def execute( 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 @@ -572,6 +1092,10 @@ async def get_node_list(self) -> list[type[IO.ComfyNode]]: return [ BriaImageEditNode, BriaRemoveImageBackground, + BriaGenFill, + BriaEraser, + BriaExpandImage, + BriaIncreaseResolution, BriaRemoveVideoBackground, BriaVideoGreenScreen, BriaVideoReplaceBackground, diff --git a/comfy_api_nodes/nodes_minimax.py b/comfy_api_nodes/nodes_minimax.py index 3c1d2925798..de3895221eb 100644 --- a/comfy_api_nodes/nodes_minimax.py +++ b/comfy_api_nodes/nodes_minimax.py @@ -3,12 +3,14 @@ import torch from typing_extensions import override -from comfy_api.latest import IO, ComfyExtension +from comfy_api.latest import IO, ComfyExtension, Input from comfy_api_nodes.apis.minimax import ( Hailuo03AudioContent, Hailuo03AudioContentUrl, + Hailuo03ContextIRRequest, Hailuo03ImageContent, Hailuo03ImageContentUrl, + Hailuo03RegenerationRequest, Hailuo03TaskCreationRequest, Hailuo03TaskCreationResponse, Hailuo03TaskQueryResponse, @@ -456,6 +458,9 @@ async def execute( 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 = [ @@ -487,10 +492,10 @@ def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = Tr IO.Int.Input( "duration", default=5, - min=5, + min=4, max=15, step=1, - tooltip="Duration of the output video in seconds (5-15).", + tooltip="Duration of the output video in seconds (4-15).", display_mode=IO.NumberDisplay.slider, ) ) @@ -939,6 +944,592 @@ async def execute( ) +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]]: @@ -950,6 +1541,8 @@ async def get_node_list(self) -> list[type[IO.ComfyNode]]: MinimaxHailuo03TextToVideoNode, MinimaxHailuo03FirstLastFrameNode, MinimaxHailuo03ReferenceNode, + MinimaxHailuo03ContextIRNode, + MinimaxHailuo03RegenerateNode, ] diff --git a/comfy_extras/nodes_minimax_music.py b/comfy_extras/nodes_minimax_music.py new file mode 100644 index 00000000000..e22103b0821 --- /dev/null +++ b/comfy_extras/nodes_minimax_music.py @@ -0,0 +1,77 @@ +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() diff --git a/nodes.py b/nodes.py index ec298e1de0b..1a3dd3f484f 100644 --- a/nodes.py +++ b/nodes.py @@ -290,6 +290,9 @@ def zero_out(self, conditioning): 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, ) @@ -1015,7 +1018,7 @@ def INPUT_TYPES(s): 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" + 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" def load_clip(self, clip_name, type="stable_diffusion", device="default"): clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) @@ -2449,6 +2452,7 @@ 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", diff --git a/requirements.txt b/requirements.txt index 4f505cc9c2f..e180e788419 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ comfyui-frontend-package==1.48.7 -comfyui-workflow-templates==0.11.40 +comfyui-workflow-templates==0.11.41 comfyui-embedded-docs==0.5.9 torch torchsde