Merge branch 'master' into fix/ltx2-prompt-empty-output
This commit is contained in:
commit
962a262d35
|
|
@ -37,7 +37,7 @@
|
|||
|
||||
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
|
||||
- ComfyUI natively supports the latest open-source state of the art models.
|
||||
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
|
||||
- [Partner nodes](https://docs.comfy.org/tutorials/partner-nodes/overview#partner-nodes) provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
|
||||
- It is available on Windows, Linux, and macOS, locally with our [desktop application](https://www.comfy.org/download), our [portable install](#installing) or on our [cloud](https://www.comfy.org/cloud).
|
||||
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
|
||||
- It integrates seamlessly into production pipelines with our API endpoints.
|
||||
|
|
|
|||
|
|
@ -180,6 +180,7 @@ parser.add_argument("--disable-async-offload", action="store_true", help="Disabl
|
|||
parser.add_argument("--disable-dynamic-vram", action="store_true", help="Disable dynamic VRAM and use estimate based model loading.")
|
||||
parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable dynamic VRAM on systems where it's not enabled by default.")
|
||||
parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.")
|
||||
parser.add_argument("--disable-cuda-graphs", action="store_true", help="Disable CUDA graphs.")
|
||||
|
||||
parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import torch
|
|||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
import comfy.model_management
|
||||
|
||||
from comfy.ldm.lightricks.model import get_timestep_embedding
|
||||
from .causal_video_autoencoder import Encoder, processor
|
||||
|
|
@ -74,8 +75,12 @@ def default_rope_dim_split(head_dim):
|
|||
|
||||
|
||||
def rope_inv_freqs(dim, base=10000.0, device=None):
|
||||
out_device = device
|
||||
if not comfy.model_management.supports_fp64(device):
|
||||
device = torch.device("cpu")
|
||||
|
||||
exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
|
||||
return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(torch.float32)
|
||||
return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(dtype=torch.float32, device=out_device)
|
||||
|
||||
|
||||
def _rope_tables(lengths, inv_freqs, device):
|
||||
|
|
|
|||
|
|
@ -91,6 +91,18 @@ def _video_t_grid(n, origin):
|
|||
return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
|
||||
|
||||
|
||||
def _ref_t_span(blk):
|
||||
# time-axis span a reference block occupies ahead of the target streams
|
||||
kind = blk["kind"]
|
||||
if kind == "image":
|
||||
return 1.0
|
||||
if kind == "audio":
|
||||
return float(blk["ref_audio_t"])
|
||||
if kind in ("video", "video_audio"):
|
||||
return max(float(blk["ref_audio_t"]), sum(_video_t_spans(blk["latent_t"])))
|
||||
return 0.0
|
||||
|
||||
|
||||
def _audio_grid(cursor, t, w_low, w_high):
|
||||
# channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0
|
||||
g = torch.zeros(t * 2, 3, dtype=torch.float64)
|
||||
|
|
@ -288,7 +300,7 @@ class FinalLayer(nn.Module):
|
|||
class PackedLayout:
|
||||
"""Static packed-sequence structure for one shape/conditioning signature."""
|
||||
|
||||
def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None):
|
||||
def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None):
|
||||
frame, w_grid = _frame_grid(latent_h, latent_w)
|
||||
frame_rows = frame.shape[0]
|
||||
|
||||
|
|
@ -299,29 +311,37 @@ class PackedLayout:
|
|||
|
||||
img_pos, img_update = [], []
|
||||
audio_pos, audio_update = [], []
|
||||
cursor = text_len
|
||||
row = text_len
|
||||
|
||||
if keyframes:
|
||||
# fl2va: keyframe cond rows right after text, sharing the target spatial grid
|
||||
for kf in keyframes:
|
||||
pixel_index = kf["resolved_frame_index"]
|
||||
if pixel_index == 0:
|
||||
cond_t = float(text_len)
|
||||
elif frame_count is not None and pixel_index == frame_count - 1:
|
||||
cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
|
||||
else:
|
||||
raise ValueError("only first/last keyframe anchors are supported")
|
||||
g = torch.empty(frame_rows, 3, dtype=torch.float64)
|
||||
g[:, 0] = cond_t
|
||||
g[:, 1:] = frame
|
||||
segments.append(("cond", frame_rows))
|
||||
pos.append(g)
|
||||
img_pos.append(torch.arange(row, row + frame_rows))
|
||||
img_update.append(torch.zeros(frame_rows, dtype=torch.bool))
|
||||
row += frame_rows
|
||||
|
||||
target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
|
||||
# refs pack between text and the targets, so the target timeline starts after their spans
|
||||
cursor = float(text_len)
|
||||
for blk in refs or ():
|
||||
cursor += _ref_t_span(blk)
|
||||
|
||||
if keyframes:
|
||||
# fl2va: keyframe cond rows right after text, sharing the target spatial grid;
|
||||
# anchors count from the target timeline origin, FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
|
||||
for kf in keyframes:
|
||||
cond_t = cursor + FRAME_RESCALE * kf["resolved_frame_index"]
|
||||
video_latent = kf.get("latent")
|
||||
if video_latent is not None:
|
||||
vt = video_latent.shape[2]
|
||||
n = vt * frame_rows
|
||||
segments.append(("cond", n))
|
||||
pos.append(_video_grid(vt, frame, cond_t))
|
||||
img_pos.append(torch.arange(row, row + n))
|
||||
img_update.append(torch.zeros(n, dtype=torch.bool))
|
||||
row += n
|
||||
audio_latent = kf.get("audio_latent")
|
||||
if audio_latent is not None:
|
||||
rt = audio_latent.shape[-1]
|
||||
segments.append(("cond_audio", rt * 2))
|
||||
pos.append(_audio_grid(cond_t, rt, *target_audio_w))
|
||||
audio_pos.append(torch.arange(row, row + rt * 2))
|
||||
audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
|
||||
row += rt * 2
|
||||
|
||||
if refs:
|
||||
cursor = float(text_len)
|
||||
for blk in refs:
|
||||
|
|
@ -389,7 +409,7 @@ class PackedLayout:
|
|||
self.audio_update = torch.cat(audio_update)
|
||||
self.signature = (text_len, latent_t, latent_h, latent_w, audio_t)
|
||||
# contiguous segment table (start, stop, kind)
|
||||
# kinds: text / cond / ref_img / ref_audio / audio / video
|
||||
# kinds: text / cond / cond_audio / ref_img / ref_audio / audio / video
|
||||
# the packed sequence is uniform per segment in (modality tag, timestep class),
|
||||
# except the text span (tag runs resolved at forward time from the presentation tags)
|
||||
seg_abs = []
|
||||
|
|
@ -529,8 +549,7 @@ class MiniMaxH3Model(nn.Module):
|
|||
if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t):
|
||||
layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t,
|
||||
keyframes=payload.get("keyframes"),
|
||||
refs=payload.get("refs"),
|
||||
frame_count=payload.get("frame_count"))
|
||||
refs=payload.get("refs"))
|
||||
|
||||
# model_base passes model_sampling.timestep(sigma) = sigma * 1000
|
||||
shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
|
||||
|
|
@ -543,14 +562,14 @@ class MiniMaxH3Model(nn.Module):
|
|||
vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
|
||||
aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
|
||||
has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments)
|
||||
has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments)
|
||||
has_aud_cond = any(k in ("cond_audio", "ref_audio") for _, _, k in layout.segments)
|
||||
seg_t = {"text": t_v, "video": t_v, "audio": t_a,
|
||||
"cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
|
||||
"ref_audio": max(t_a, aud_aug)}
|
||||
"cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)}
|
||||
unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set())
|
||||
| ({seg_t["ref_audio"]} if has_aud_cond else set()))
|
||||
t_row = {t: i for i, t in enumerate(unique_t)}
|
||||
seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
|
||||
seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2}
|
||||
|
||||
text_tags = payload.get("text_token_tags")
|
||||
mod_segments = []
|
||||
|
|
|
|||
|
|
@ -0,0 +1,343 @@
|
|||
import dataclasses
|
||||
import hashlib
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
from comfy.text_encoders.llama import Llama2_, Qwen3_8BConfig
|
||||
|
||||
from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS
|
||||
|
||||
|
||||
CFG_SCALE = 1.5
|
||||
CFG_TOP_K = 50
|
||||
C0_VOCAB_SIZE = 16384
|
||||
MAX_PROMPT_TOKENS = 5000
|
||||
MAX_AUDIO_FRAMES = 9000
|
||||
AUDIO_FRAMES_PER_SECOND = 25
|
||||
|
||||
|
||||
def derive_seed(seed, *parts):
|
||||
digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm")
|
||||
digest.update(int(seed).to_bytes(8, "little", signed=False))
|
||||
for part in parts:
|
||||
value = str(part).encode("utf-8")
|
||||
digest.update(len(value).to_bytes(4, "little"))
|
||||
digest.update(value)
|
||||
return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1)
|
||||
|
||||
|
||||
def sample_topk(logits, top_k, generator):
|
||||
values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9)
|
||||
top_k = min(top_k, values.shape[-1])
|
||||
threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None]
|
||||
values = values.masked_fill(values < threshold, -float("inf"))
|
||||
probabilities = torch.nan_to_num(torch.softmax(values, dim=-1), nan=0.0)
|
||||
probabilities = probabilities / probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12)
|
||||
return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1)
|
||||
|
||||
|
||||
class RVQAttention(nn.Module):
|
||||
def __init__(self, hidden_size, num_heads, merged_qkv, dtype, device, operations):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
self.merged_qkv = merged_qkv
|
||||
if merged_qkv:
|
||||
self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
|
||||
else:
|
||||
self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
batch, length, hidden_size = x.shape
|
||||
if self.merged_qkv:
|
||||
q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
|
||||
else:
|
||||
q = self.q_proj(x)
|
||||
k = self.k_proj(x)
|
||||
v = self.v_proj(x)
|
||||
q = q.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
k = k.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
v = v.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
mask = torch.full((length, length), torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype).triu_(1)
|
||||
attention = optimized_attention_for_device(q.device, mask=True, small_input=True)
|
||||
out = attention(q, k, v, self.num_heads, mask=mask, skip_reshape=True)
|
||||
return self.o_proj(out)
|
||||
|
||||
|
||||
class RVQRMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, dtype, device):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device))
|
||||
|
||||
def forward(self, x):
|
||||
return torch.nn.functional.rms_norm(x, (x.shape[-1],), comfy.ops.cast_to_input(self.weight, x), 1e-6)
|
||||
|
||||
|
||||
class RVQMLP(nn.Module):
|
||||
def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations):
|
||||
super().__init__()
|
||||
self.merged_mlp = merged_mlp
|
||||
if merged_mlp:
|
||||
self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
|
||||
else:
|
||||
self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
|
||||
self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
|
||||
self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
if self.merged_mlp:
|
||||
return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu")
|
||||
return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class RVQDecoderBlock(nn.Module):
|
||||
def __init__(self, hidden_size, num_heads, intermediate_size, merged_qkv, merged_mlp, dtype, device, operations):
|
||||
super().__init__()
|
||||
self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device)
|
||||
self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations)
|
||||
self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device)
|
||||
self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, dtype, device, operations)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.self_attn(self.input_layernorm(x))
|
||||
return x + self.mlp(self.post_attention_layernorm(x))
|
||||
|
||||
|
||||
class RVQDepthDecoder(nn.Module):
|
||||
def __init__(self, config, dtype, device, operations):
|
||||
super().__init__()
|
||||
hidden_size = int(config["hidden_size"])
|
||||
audio_vocab_size = int(config["audio_vocab_size"])
|
||||
merged_qkv = config.get("decoder_merged_qkv", False)
|
||||
merged_mlp = config.get("decoder_merged_mlp", False)
|
||||
num_codebooks = int(config["audio_num_codebooks"])
|
||||
self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
|
||||
self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device)
|
||||
self.audio_heads = nn.ModuleList([
|
||||
operations.Linear(hidden_size, audio_vocab_size, bias=False, dtype=dtype, device=device)
|
||||
for _ in range(num_codebooks - 1)
|
||||
])
|
||||
self.layers = nn.ModuleList([
|
||||
RVQDecoderBlock(
|
||||
hidden_size,
|
||||
int(config["decoder_num_heads"]),
|
||||
int(config["decoder_intermediate_size"]),
|
||||
merged_qkv,
|
||||
merged_mlp,
|
||||
dtype,
|
||||
device,
|
||||
operations,
|
||||
)
|
||||
for _ in range(int(config["decoder_num_layers"]))
|
||||
])
|
||||
self.norm = RVQRMSNorm(hidden_size, dtype, device)
|
||||
|
||||
def forward(self, sequence):
|
||||
positions = torch.arange(sequence.shape[1], device=sequence.device)
|
||||
x = sequence + self.pos_embedding(positions, out_dtype=sequence.dtype).unsqueeze(0)
|
||||
for layer in self.layers:
|
||||
x = layer(x)
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
class MiniMaxMusic3AR(nn.Module):
|
||||
def __init__(self, config, dtype, device, operations):
|
||||
super().__init__()
|
||||
config_fields = {field.name for field in dataclasses.fields(Qwen3_8BConfig)}
|
||||
qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields})
|
||||
qwen_config.lm_head = False
|
||||
qwen_config.fixed_kv = True
|
||||
self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations)
|
||||
self.model.prefetch_dynamic_vbars = True
|
||||
self.model.graph_dynamic_vbar_blocks = True
|
||||
self.model.lm_head = operations.Linear(qwen_config.hidden_size, qwen_config.vocab_size, bias=False, dtype=dtype, device=device)
|
||||
self.model.lm_head_pruned = operations.Linear(qwen_config.hidden_size, C0_VOCAB_SIZE + 1, bias=False, dtype=dtype, device=device)
|
||||
self.model.embed_tokens_prefill = operations.Embedding(AUDIO_CODE_OFFSET, qwen_config.hidden_size, dtype=dtype, device=device)
|
||||
self.model.embed_tokens_audio = operations.Embedding(C0_VOCAB_SIZE, qwen_config.hidden_size, dtype=dtype, device=device)
|
||||
self.model.pruned_lm_head = None
|
||||
self.model.pruned_embedding = None
|
||||
self.model.audio_extra_embedding = operations.Embedding(
|
||||
int(config["audio_vocab_size"]) * (int(config["audio_num_codebooks"]) - 1),
|
||||
qwen_config.hidden_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
self.model.audio_decoder = RVQDepthDecoder(config, dtype, device, operations)
|
||||
self.audio_vocab_size = int(config["audio_vocab_size"])
|
||||
self.num_codebooks = int(config["audio_num_codebooks"])
|
||||
self.embedding_scale = self.num_codebooks ** -0.5
|
||||
|
||||
def _guided_c0(self, logits, cfg_scale, top_k):
|
||||
conditioned = logits[0:1].float()
|
||||
unconditioned = logits[1:2].float()
|
||||
guided = unconditioned + (conditioned - unconditioned) * cfg_scale
|
||||
threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None]
|
||||
return guided.masked_fill(conditioned < threshold, -float("inf"))
|
||||
|
||||
def _depth_codes(self, hidden, c0, c0_embed, generator, execution_dtype, cfg_scale, top_k):
|
||||
decoder = self.model.audio_decoder
|
||||
sequence = [decoder.projection(hidden).unsqueeze(1)]
|
||||
sequence.append(decoder.projection(c0_embed).unsqueeze(1))
|
||||
codes = [c0]
|
||||
hidden_parts = []
|
||||
for index in range(1, self.num_codebooks):
|
||||
out = decoder(torch.cat(sequence, dim=1))[:, -1]
|
||||
hidden_parts.append(out[:1].detach())
|
||||
logits = decoder.audio_heads[index - 1](out)
|
||||
conditioned = logits[:1].float()
|
||||
unconditioned = logits[1:2].float()
|
||||
code = sample_topk(unconditioned + (conditioned - unconditioned) * cfg_scale, top_k, generator).repeat(2)
|
||||
codes.append(code)
|
||||
if index < self.num_codebooks - 1:
|
||||
embedding = self.model.audio_extra_embedding(
|
||||
code + (index - 1) * self.audio_vocab_size,
|
||||
out_dtype=execution_dtype,
|
||||
)
|
||||
sequence.append(decoder.projection(embedding).unsqueeze(1))
|
||||
return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1)
|
||||
|
||||
def _embed_c0(self, codes, execution_dtype):
|
||||
if self.model.pruned_embedding:
|
||||
return self.model.embed_tokens_audio(codes, out_dtype=execution_dtype)
|
||||
return self.model.embed_tokens(codes + AUDIO_CODE_OFFSET, out_dtype=execution_dtype)
|
||||
|
||||
def _embed_audio_frame(self, codes, execution_dtype):
|
||||
c0 = self._embed_c0(codes[:, 0], execution_dtype)
|
||||
offsets = torch.arange(self.num_codebooks - 1, device=codes.device) * self.audio_vocab_size
|
||||
extra = self.model.audio_extra_embedding(codes[:, 1:] + offsets.unsqueeze(0), out_dtype=execution_dtype).sum(dim=1)
|
||||
return ((c0 + extra) * self.embedding_scale).unsqueeze(1)
|
||||
|
||||
def _sample_c0(self, hidden, cfg_scale, top_k, generator, vocab_mask):
|
||||
if self.model.pruned_lm_head:
|
||||
guided = self._guided_c0(self.model.lm_head_pruned(hidden).float(), cfg_scale, top_k)
|
||||
code = sample_topk(guided, top_k, generator)
|
||||
stop_token = 0
|
||||
offset = 1
|
||||
else:
|
||||
logits = self.model.lm_head(hidden).float()
|
||||
stop_token = SPECIAL_TOKEN_IDS["<|audio_end|>"]
|
||||
logits = logits.masked_fill(vocab_mask, -float("inf"))
|
||||
guided = self._guided_c0(logits, cfg_scale, top_k).masked_fill(vocab_mask, -float("inf"))
|
||||
code = sample_topk(guided, top_k, generator)
|
||||
offset = AUDIO_CODE_OFFSET
|
||||
return torch.where(code == stop_token, 0, code - offset), code, stop_token
|
||||
|
||||
def generate(self, input_ids, seed, max_audio_frames, device, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K):
|
||||
prompt_tokens = int(input_ids.shape[1])
|
||||
if prompt_tokens > MAX_PROMPT_TOKENS:
|
||||
raise ValueError(f"MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is {MAX_PROMPT_TOKENS}")
|
||||
|
||||
input_ids = input_ids.to(device)
|
||||
if comfy.model_management.should_use_bf16(device):
|
||||
execution_dtype = torch.bfloat16
|
||||
else:
|
||||
execution_dtype = torch.float32
|
||||
unconditioned = input_ids.clone()
|
||||
unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"]
|
||||
text_ids = torch.cat((input_ids, unconditioned), dim=0)
|
||||
if self.model.pruned_embedding:
|
||||
text_embeds = self.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype)
|
||||
else:
|
||||
text_embeds = self.model.embed_tokens(text_ids, out_dtype=execution_dtype)
|
||||
decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES)
|
||||
past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype)
|
||||
output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype)
|
||||
last_hidden = output[0][:, -1]
|
||||
past = output[2]
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar"))
|
||||
decoder = self.model.audio_decoder
|
||||
depth_io = {
|
||||
"hidden": torch.empty_like(last_hidden),
|
||||
"c0": torch.empty((last_hidden.shape[0],), dtype=torch.long, device=device),
|
||||
"c0_embed": torch.empty_like(last_hidden),
|
||||
"codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device),
|
||||
"depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device),
|
||||
}
|
||||
decoder._comfy_cross_step_state = depth_io
|
||||
comfy.model_management._register_cross_step(decoder)
|
||||
hidden_frames = []
|
||||
pending_code = None
|
||||
stop_token = None
|
||||
pending_event = None
|
||||
pending_hidden = None
|
||||
progress = comfy.utils.ProgressBar(decode_limit)
|
||||
cuda_device = torch.device(device).type == "cuda"
|
||||
vocab_mask = None
|
||||
if not self.model.pruned_lm_head:
|
||||
vocab_mask = torch.ones(self.model.vocab_size, dtype=torch.bool, device=device)
|
||||
vocab_mask[AUDIO_CODE_OFFSET:AUDIO_CODE_OFFSET + C0_VOCAB_SIZE] = False
|
||||
vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False
|
||||
|
||||
for frame_index in comfy.utils.model_trange(decode_limit + 1, desc="AR sampling"):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
if pending_code is not None:
|
||||
if pending_event is not None:
|
||||
pending_event.synchronize()
|
||||
if int(pending_code.item()) == stop_token:
|
||||
pending_hidden = None
|
||||
break
|
||||
if pending_hidden is not None:
|
||||
hidden_frames.append(pending_hidden)
|
||||
progress.update_absolute(len(hidden_frames))
|
||||
if len(hidden_frames) >= decode_limit:
|
||||
break
|
||||
|
||||
c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
|
||||
if pending_code is None:
|
||||
pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device)
|
||||
if cuda_device:
|
||||
pending_event = torch.cuda.Event()
|
||||
pending_code.copy_(code_or_stop, non_blocking=cuda_device)
|
||||
if pending_event is not None:
|
||||
pending_event.record()
|
||||
|
||||
c0 = c0.repeat(2)
|
||||
c0_embed = self._embed_c0(c0, execution_dtype)
|
||||
depth_io["hidden"].copy_(last_hidden)
|
||||
depth_io["c0"].copy_(c0)
|
||||
depth_io["c0_embed"].copy_(c0_embed)
|
||||
|
||||
def depth_core():
|
||||
codes, depth_hidden = self._depth_codes(
|
||||
depth_io["hidden"], depth_io["c0"], depth_io["c0_embed"], generator, execution_dtype, cfg_scale, top_k
|
||||
)
|
||||
depth_io["codes"].copy_(codes)
|
||||
depth_io["depth_hidden"].copy_(depth_hidden)
|
||||
|
||||
depth_queue = comfy.model_prefetch.make_prefetch_queue(
|
||||
[[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True}
|
||||
)
|
||||
comfy.model_prefetch.prefetch_queue_pop(
|
||||
depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator
|
||||
)
|
||||
comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None)
|
||||
feedback_codes = depth_io["codes"]
|
||||
depth_hidden = depth_io["depth_hidden"]
|
||||
frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1)
|
||||
if frame_index > 0:
|
||||
pending_hidden = frame_hidden[0].clone()
|
||||
|
||||
feedback = self._embed_audio_frame(feedback_codes, execution_dtype)
|
||||
output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
|
||||
last_hidden = output[0][:, -1]
|
||||
past = output[2]
|
||||
|
||||
if pending_hidden is not None and len(hidden_frames) < decode_limit:
|
||||
if pending_event is not None:
|
||||
pending_event.synchronize()
|
||||
if int(pending_code.item()) != stop_token:
|
||||
hidden_frames.append(pending_hidden)
|
||||
|
||||
if not hidden_frames:
|
||||
raise ValueError("MiniMax Music3 generated zero audio frames")
|
||||
return torch.stack(hidden_frames).to(device="cpu")
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
|
||||
lines.append(line.rstrip())
|
||||
text = "\n".join(lines)
|
||||
text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
|
||||
return text.replace("• ", "").replace(" ", "")
|
||||
|
||||
|
||||
def clean_caption(caption):
|
||||
def replace_special(match):
|
||||
inner = match.group(1).strip()
|
||||
parts = inner.split(None, 1)
|
||||
return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner
|
||||
|
||||
text = _SPECIAL_TAG_RE.sub(replace_special, caption)
|
||||
text = _remove_markdown_format(text)
|
||||
return re.sub(r"\n{2,}", "\n", text)
|
||||
|
||||
|
||||
def normalize_lyrics(lyrics):
|
||||
parts = _LYRIC_TAG_RE.split(lyrics)
|
||||
text = "\n".join(part.lower() if part.startswith("[") else part for part in parts if part)
|
||||
text = text.replace(" ^ ", "\n")
|
||||
return f"[start]\n{text}"
|
||||
|
||||
|
||||
def build_prompt(caption, lyrics):
|
||||
return (
|
||||
"<|im_start|><|caption_start|>"
|
||||
f"{clean_caption(caption)}"
|
||||
"<|caption_end|><|lyrics_start|>"
|
||||
f"{normalize_lyrics(lyrics)}"
|
||||
"<|lyrics_end|><|im_end|><|audio_start|>"
|
||||
)
|
||||
|
||||
|
||||
def validate_tokenizer(tokenizer):
|
||||
for token, expected in SPECIAL_TOKEN_IDS.items():
|
||||
token_id = tokenizer.convert_tokens_to_ids(token)
|
||||
if token_id != expected:
|
||||
raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}: expected {expected}, got {token_id}")
|
||||
|
|
@ -22,6 +22,7 @@ import torch
|
|||
import logging
|
||||
import comfy.ldm.lightricks.av_model
|
||||
import comfy.ldm.minimax.model
|
||||
import comfy.ldm.minimax_music.dit
|
||||
import comfy.nested_tensor
|
||||
import comfy.ldm.lightricks.symmetric_patchifier
|
||||
import comfy.context_windows
|
||||
|
|
@ -2164,13 +2165,13 @@ class MiniMaxH3(BaseModel):
|
|||
keyframes = kwargs.get("minimax_keyframes", None)
|
||||
if keyframes is not None:
|
||||
payload["keyframes"] = keyframes
|
||||
payload["frame_count"] = kwargs.get("minimax_frame_count", None)
|
||||
payload["cond_video_latents"] = [kf["latent"] for kf in keyframes]
|
||||
payload["cond_video_latents"] = [kf["latent"] for kf in keyframes if kf.get("latent") is not None]
|
||||
payload["cond_audio_latents"] = [kf["audio_latent"] for kf in keyframes if kf.get("audio_latent") is not None]
|
||||
refs = kwargs.get("minimax_refs", None)
|
||||
if refs is not None:
|
||||
payload["refs"] = refs
|
||||
payload["cond_video_latents"] = [r["latent"] for r in refs if "latent" in r]
|
||||
payload["cond_audio_latents"] = [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
|
||||
payload["cond_video_latents"] = payload.get("cond_video_latents", []) + [r["latent"] for r in refs if "latent" in r]
|
||||
payload["cond_audio_latents"] = payload.get("cond_audio_latents", []) + [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
|
||||
if kwargs.get("minimax_visual_cond_noise_aug", None) is not None:
|
||||
payload["visual_cond_noise_aug"] = kwargs["minimax_visual_cond_noise_aug"]
|
||||
if kwargs.get("minimax_audio_cond_noise_aug", None) is not None:
|
||||
|
|
@ -2184,7 +2185,7 @@ class MiniMaxH3(BaseModel):
|
|||
payload["layout"] = comfy.ldm.minimax.model.PackedLayout(
|
||||
cross_attn.shape[1], vs[2], (vs[3] + 1) // 2 * 2, (vs[4] + 1) // 2 * 2,
|
||||
latent_shapes[1][-1], keyframes=payload.get("keyframes"),
|
||||
refs=payload.get("refs"), frame_count=payload.get("frame_count"))
|
||||
refs=payload.get("refs"))
|
||||
out['minimax_payload'] = comfy.conds.CONDConstant(payload)
|
||||
return out
|
||||
|
||||
|
|
@ -2337,6 +2338,18 @@ class ACEStep15(BaseModel):
|
|||
out['refer_audio'] = comfy.conds.CONDRegular(refer_audio)
|
||||
return out
|
||||
|
||||
class MiniMaxMusic3(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax_music.dit.MiniMaxMusic3DiT)
|
||||
|
||||
def process_timestep(self, timestep, **kwargs):
|
||||
return 1.0 - timestep
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
out["conditioning_scale"] = comfy.conds.CONDRegular(kwargs["conditioning_scale"])
|
||||
return out
|
||||
|
||||
class Omnigen2(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.omnigen.omnigen2.OmniGen2Transformer2DModel)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
@ -830,11 +837,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
|||
|
||||
dit_config["use_adaln_lora"] = True
|
||||
dit_config["adaln_lora_dim"] = 256
|
||||
dit_config["num_blocks"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
|
||||
if dit_config["model_channels"] == 2048:
|
||||
dit_config["num_blocks"] = 28
|
||||
dit_config["num_heads"] = 16
|
||||
elif dit_config["model_channels"] == 5120:
|
||||
dit_config["num_blocks"] = 36
|
||||
dit_config["num_heads"] = 40
|
||||
|
||||
if dit_config["in_channels"] == 16:
|
||||
|
|
|
|||
|
|
@ -490,28 +490,36 @@ try:
|
|||
except:
|
||||
rocm_version = (6, -1)
|
||||
|
||||
def aotriton_supported(gpu_arch):
|
||||
path = torch.__path__[0]
|
||||
path = os.path.join(os.path.join(path, "lib"), "aotriton.images")
|
||||
gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path))))
|
||||
if gpu_arch in gfx:
|
||||
return True
|
||||
if "{}x".format(gpu_arch[:-1]) in gfx:
|
||||
return True
|
||||
if "{}xx".format(gpu_arch[:-2]) in gfx:
|
||||
return True
|
||||
return False
|
||||
def aotriton_supported():
|
||||
"""Whether pytorch reports flash attention as usable on this gpu.
|
||||
|
||||
can_use_flash_attention() evaluates runtime eligibility for the given
|
||||
parameters; on a ROCm build that includes checking the gpu arch against the
|
||||
kernel images AOTriton was compiled for. Querying it avoids assuming where
|
||||
those images live inside the torch install. The probe tensor is shaped and
|
||||
typed to pass the unrelated SDPA checks, so False means no hardware support
|
||||
rather than a rejected shape.
|
||||
"""
|
||||
try:
|
||||
if not torch.backends.cuda.is_flash_attention_available(): # not built with flash attention
|
||||
return False
|
||||
q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
|
||||
params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
|
||||
return torch.backends.cuda.can_use_flash_attention(params, False)
|
||||
except (AttributeError, RuntimeError, TypeError) as e:
|
||||
logging.warning("Could not query aotriton support: {}".format(e))
|
||||
return False
|
||||
|
||||
logging.info("AMD arch: {}".format(arch))
|
||||
logging.info("ROCm version: {}".format(rocm_version))
|
||||
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
|
||||
if aotriton_supported(arch): # AMD efficient attention implementation depends on aotriton.
|
||||
if aotriton_supported(): # AMD efficient attention implementation depends on aotriton.
|
||||
if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much
|
||||
if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]): # TODO: more arches, TODO: gfx950
|
||||
ENABLE_PYTORCH_ATTENTION = True
|
||||
if rocm_version >= (7, 0):
|
||||
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
|
||||
ENABLE_PYTORCH_ATTENTION = True
|
||||
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
|
||||
ENABLE_PYTORCH_ATTENTION = True
|
||||
if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
|
||||
if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
|
||||
SUPPORT_FP8_OPS = True
|
||||
|
|
@ -1360,9 +1368,14 @@ STREAM_CAST_BUFFERS = {}
|
|||
LARGEST_CASTED_WEIGHT = (None, 0)
|
||||
STREAM_AIMDO_CAST_BUFFERS = {}
|
||||
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
|
||||
CROSS_STEP_STATE = weakref.WeakSet()
|
||||
|
||||
DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
|
||||
|
||||
# NOTE: devs/agents: this is temporary and will be removed in a future comfy. Not supported for custom node use.
|
||||
def _register_cross_step(module):
|
||||
CROSS_STEP_STATE.add(module)
|
||||
|
||||
def get_cast_buffer(offload_stream, device, size, ref):
|
||||
global LARGEST_CASTED_WEIGHT
|
||||
|
||||
|
|
@ -1417,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():
|
||||
|
|
|
|||
|
|
@ -1887,8 +1887,29 @@ class ModelPatcherDynamic(ModelPatcher):
|
|||
loading = self._load_list(for_dynamic=True, default_device=device_to)
|
||||
loading.sort()
|
||||
|
||||
get_units = getattr(self.model, "get_dynamic_vram__units", None)
|
||||
dynamic_units, last_dynamic_units = get_units() if get_units is not None else ([], [])
|
||||
dynamic_units = list(dynamic_units)
|
||||
last_dynamic_units = list(last_dynamic_units)
|
||||
loading_by_module = {entry[-2]: entry for entry in loading}
|
||||
loading = []
|
||||
for unit in dynamic_units:
|
||||
unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
|
||||
modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
|
||||
for index, module in enumerate(modules):
|
||||
loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
|
||||
last_loading = []
|
||||
for unit in last_dynamic_units:
|
||||
unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
|
||||
modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
|
||||
for index, module in enumerate(modules):
|
||||
last_loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
|
||||
loading.extend((*entry, None) for entry in loading_by_module.values())
|
||||
loading.extend(last_loading)
|
||||
v_block = None
|
||||
|
||||
for x in loading:
|
||||
*_, module_mem, n, m, params = 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 @@ class ModelPatcherDynamic(ModelPatcher):
|
|||
|
||||
move_weight_functions(m, device_to)
|
||||
|
||||
if hasattr(m, "_v"):
|
||||
v_block = m._v if v_block is None else (v_block[0], v_block[1], max(v_block[2], m._v[1] + m._v[2] - v_block[1]))
|
||||
if end_of_block is not None:
|
||||
unit = end_of_block
|
||||
(unit[0] if isinstance(unit, (list, tuple)) else unit)._v_block = v_block
|
||||
v_block = None
|
||||
|
||||
for key, buf in self.model.named_buffers(recurse=True):
|
||||
if key not in self.backup_buffers:
|
||||
self.backup_buffers[key] = buf
|
||||
|
|
|
|||
|
|
@ -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) and getattr(module, "_v_block", None) is not None
|
||||
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)
|
||||
|
|
|
|||
14
comfy/ops.py
14
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 cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
|
|||
for s in comfy_modules:
|
||||
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
|
||||
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
|
||||
if return_faulted and (signature is None or not resident):
|
||||
fully_faulted = False
|
||||
prefetch = {
|
||||
"signature": signature,
|
||||
"resident": resident,
|
||||
|
|
@ -255,10 +259,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
|
|||
prefetch["needs_cast"] = needs_cast
|
||||
s._prefetch = prefetch
|
||||
|
||||
if return_faulted:
|
||||
return offload_stream, fully_faulted
|
||||
return offload_stream
|
||||
|
||||
|
||||
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant):
|
||||
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 resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
|
|||
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 resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
|
|||
if prefetch["signature"] is not None:
|
||||
prefetch["resident"] = True
|
||||
|
||||
return weight, bias
|
||||
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):
|
||||
|
|
|
|||
31
comfy/sd.py
31
comfy/sd.py
|
|
@ -25,6 +25,7 @@ import comfy.ldm.cogvideo.vae
|
|||
import comfy.ldm.hunyuan_video.vae
|
||||
import comfy.ldm.mmaudio.vae.autoencoder
|
||||
import comfy.ldm.audio.vae_sa3
|
||||
import comfy.ldm.minimax_music.dav
|
||||
import comfy.pixel_space_convert
|
||||
import comfy.weight_adapter
|
||||
import yaml
|
||||
|
|
@ -32,6 +33,7 @@ import math
|
|||
import os
|
||||
|
||||
import comfy.utils
|
||||
import comfy.ops
|
||||
|
||||
from . import clip_vision
|
||||
from . import gligen
|
||||
|
|
@ -74,6 +76,7 @@ import comfy.text_encoders.longcat_image
|
|||
import comfy.text_encoders.qwen35
|
||||
import comfy.text_encoders.qwen3vl
|
||||
import comfy.text_encoders.minimax
|
||||
import comfy.text_encoders.minimax_music
|
||||
import comfy.ldm.minimax.vae
|
||||
import comfy.ldm.minimax.audio_vae
|
||||
import comfy.text_encoders.boogu
|
||||
|
|
@ -515,7 +518,22 @@ class VAE:
|
|||
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,16 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
|||
clip_target.params = {}
|
||||
if len(clip_data) == 1:
|
||||
te_model = detect_te_model(clip_data[0])
|
||||
if 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.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0])
|
||||
clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel
|
||||
clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer
|
||||
elif te_model == TEModel.CLIP_G:
|
||||
if clip_type == CLIPType.STABLE_CASCADE:
|
||||
clip_target.clip = sdxl_clip.StableCascadeClipModel
|
||||
clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import comfy.text_encoders.genmo
|
|||
import comfy.text_encoders.lt
|
||||
import comfy.text_encoders.hunyuan_video
|
||||
import comfy.text_encoders.minimax
|
||||
import comfy.text_encoders.minimax_music
|
||||
import comfy.text_encoders.cosmos
|
||||
import comfy.text_encoders.lumina2
|
||||
import comfy.text_encoders.wan
|
||||
|
|
@ -2200,6 +2201,28 @@ class ACEStep15(supported_models_base.BASE):
|
|||
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
|
||||
|
||||
class MiniMaxMusic3(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"audio_model": "minimax_music3",
|
||||
}
|
||||
|
||||
latent_format = comfy.latent_formats.MiniMaxMusic3
|
||||
memory_usage_factor = 2.0
|
||||
supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
sampling_settings = {"multiplier": 1.0}
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
return model_base.MiniMaxMusic3(self, device=device)
|
||||
|
||||
def model_type(self, state_dict, prefix=""):
|
||||
return model_base.ModelType.FLOW
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0])
|
||||
target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
|
||||
target.params["projection_config"] = detect
|
||||
return target
|
||||
|
||||
|
||||
class LongCatImage(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
|
|
@ -2494,6 +2517,7 @@ models = [
|
|||
ChromaRadiance,
|
||||
ACEStep,
|
||||
ACEStep15,
|
||||
MiniMaxMusic3,
|
||||
Omnigen2,
|
||||
Boogu,
|
||||
MageFlow,
|
||||
|
|
|
|||
|
|
@ -1184,6 +1184,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po
|
|||
|
||||
class Gemma4_Tokenizer():
|
||||
tokenizer_json_data = None
|
||||
prime_empty_thought = False
|
||||
|
||||
def state_dict(self):
|
||||
if self.tokenizer_json_data is not None:
|
||||
|
|
@ -1334,8 +1335,8 @@ class Gemma4_Tokenizer():
|
|||
num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
|
||||
n_audio_tokens = self._audio_token_count(num_samples)
|
||||
media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
|
||||
# Non-thinking mode primes an empty thought channel so the model answers directly.
|
||||
model_open = "" if thinking else "<|channel>thought\n<channel|>"
|
||||
# 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline.
|
||||
model_open = "<|channel>thought\n<channel|>" if self.prime_empty_thought and not thinking else ""
|
||||
llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
|
||||
|
||||
text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
|
||||
|
|
@ -1421,6 +1422,7 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
|
|||
class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
|
||||
"""Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
|
||||
embedding_size = 3840
|
||||
prime_empty_thought = True
|
||||
|
||||
def _extract_audio_features(self, waveform, sample_rate):
|
||||
audio = self._resample_16k(waveform, sample_rate)
|
||||
|
|
@ -1503,7 +1505,7 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
|
|||
|
||||
# Variants
|
||||
|
||||
def _make_variant(config_cls):
|
||||
def _make_variant(config_cls, prime_empty_thought=False):
|
||||
audio = config_cls.audio_config is not None
|
||||
bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,)
|
||||
class Variant(*bases):
|
||||
|
|
@ -1513,8 +1515,8 @@ def _make_variant(config_cls):
|
|||
if audio:
|
||||
self._init_audio(self.model.config, dtype, device, operations)
|
||||
embedding_size = config_cls.hidden_size
|
||||
if embedding_size != Gemma4SDTokenizer.embedding_size:
|
||||
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size})
|
||||
if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought:
|
||||
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought})
|
||||
class Tokenizer(Gemma4Tokenizer):
|
||||
tokenizer_class = tok_cls
|
||||
Variant.tokenizer = Tokenizer
|
||||
|
|
@ -1524,7 +1526,7 @@ def _make_variant(config_cls):
|
|||
|
||||
Gemma4_E4B = _make_variant(Gemma4Config)
|
||||
Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
|
||||
Gemma4_31B = _make_variant(Gemma4_31B_Config)
|
||||
Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True)
|
||||
|
||||
|
||||
# Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).
|
||||
|
|
|
|||
|
|
@ -5,15 +5,33 @@ from typing import Optional, Any, Tuple
|
|||
import math
|
||||
from tqdm import tqdm
|
||||
import comfy.utils
|
||||
import comfy_kitchen
|
||||
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.ops
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.clip_model
|
||||
|
||||
from . import qwen_vl
|
||||
|
||||
|
||||
@dataclass
|
||||
class FixedKV:
|
||||
key: torch.Tensor
|
||||
value: torch.Tensor
|
||||
index: int
|
||||
position: torch.Tensor
|
||||
seqlen: torch.Tensor
|
||||
|
||||
def prepare(self, num_tokens):
|
||||
self.position.fill_(self.index)
|
||||
self.seqlen.fill_(self.index + num_tokens)
|
||||
|
||||
def advance(self, num_tokens):
|
||||
self.index += num_tokens
|
||||
|
||||
@dataclass
|
||||
class Llama2Config:
|
||||
vocab_size: int = 128320
|
||||
|
|
@ -249,6 +267,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 +519,14 @@ class Attention(nn.Module):
|
|||
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:
|
||||
self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype)
|
||||
else:
|
||||
self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
|
||||
self.k_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
|
||||
self.v_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
|
||||
self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype)
|
||||
|
||||
self.q_norm = None
|
||||
|
|
@ -522,9 +548,12 @@ class Attention(nn.Module):
|
|||
):
|
||||
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 +566,29 @@ class Attention(nn.Module):
|
|||
|
||||
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 +619,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:
|
||||
self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype)
|
||||
else:
|
||||
self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.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 +658,7 @@ class TransformerBlock(nn.Module):
|
|||
optimized_attention=None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
):
|
||||
output = x
|
||||
# Self Attention
|
||||
residual = x
|
||||
x = self.input_layernorm(x)
|
||||
|
|
@ -612,7 +675,7 @@ class TransformerBlock(nn.Module):
|
|||
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 +704,7 @@ class TransformerBlockGemma2(nn.Module):
|
|||
optimized_attention=None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
):
|
||||
output = x
|
||||
sliding_window = None
|
||||
if self.transformer_type == 'gemma3':
|
||||
if self.sliding_attention:
|
||||
|
|
@ -676,7 +740,7 @@ class TransformerBlockGemma2(nn.Module):
|
|||
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 +752,14 @@ def _make_scaled_embedding(ops, vocab_size, hidden_size, scale, device, dtype):
|
|||
|
||||
|
||||
class Llama2_(nn.Module):
|
||||
fixed_kv = False
|
||||
graph_dynamic_vbar_blocks = False
|
||||
|
||||
def __init__(self, config, device=None, dtype=None, ops=None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.fixed_kv = getattr(config, "fixed_kv", False)
|
||||
self.graph_dynamic_vbar_blocks = False
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
|
||||
|
|
@ -713,8 +782,27 @@ class Llama2_(nn.Module):
|
|||
if config.lm_head:
|
||||
self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
|
||||
|
||||
def get_dynamic_vram__units(self):
|
||||
return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
|
||||
|
||||
def get_past_len(self, past_key_values):
|
||||
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 +844,33 @@ class Llama2_(nn.Module):
|
|||
|
||||
optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
|
||||
|
||||
fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV)
|
||||
enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv and seq_len == 1 and mask is None
|
||||
if enable_graph:
|
||||
freqs_cis_groups = freqs_cis if isinstance(freqs_cis, list) else [freqs_cis]
|
||||
cross_step_state_key = [(x.shape, x.stride(), x.dtype, x.device)]
|
||||
for group in freqs_cis_groups:
|
||||
for tensor in group:
|
||||
cross_step_state_key.append((tensor.shape, tensor.stride(), tensor.dtype, tensor.device))
|
||||
cross_step_state_key = tuple(cross_step_state_key)
|
||||
cross_step_state = getattr(self, "_comfy_cross_step_state", None)
|
||||
if cross_step_state is None or cross_step_state["key"] != cross_step_state_key:
|
||||
static_freqs_cis = []
|
||||
for group in freqs_cis_groups:
|
||||
static_freqs_cis.append(tuple(torch.empty_like(tensor) for tensor in group))
|
||||
if not isinstance(freqs_cis, list):
|
||||
static_freqs_cis = static_freqs_cis[0]
|
||||
cross_step_state = {"key": cross_step_state_key, "x": torch.empty_like(x), "freqs_cis": static_freqs_cis}
|
||||
self._comfy_cross_step_state = cross_step_state
|
||||
comfy.model_management._register_cross_step(self)
|
||||
cross_step_state["x"].copy_(x)
|
||||
static_freqs_cis_groups = cross_step_state["freqs_cis"] if isinstance(freqs_cis, list) else [cross_step_state["freqs_cis"]]
|
||||
for source_group, target_group in zip(freqs_cis_groups, static_freqs_cis_groups):
|
||||
for source, target in zip(source_group, target_group):
|
||||
target.copy_(source)
|
||||
x = cross_step_state["x"]
|
||||
freqs_cis = cross_step_state["freqs_cis"]
|
||||
|
||||
intermediate = None
|
||||
all_intermediate = None
|
||||
only_layers = None
|
||||
|
|
@ -769,7 +884,8 @@ class Llama2_(nn.Module):
|
|||
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 +895,24 @@ class Llama2_(nn.Module):
|
|||
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 fixed_kv:
|
||||
past_kv.prepare(seq_len)
|
||||
|
||||
if current_kv is not None:
|
||||
next_key_values.append(current_kv)
|
||||
def core():
|
||||
nonlocal x
|
||||
x, current_kv = layer(
|
||||
x=x,
|
||||
attention_mask=mask,
|
||||
freqs_cis=freqs_cis,
|
||||
optimized_attention=optimized_attention,
|
||||
past_key_value=past_kv,
|
||||
)
|
||||
if next_key_values:
|
||||
next_key_values[i] = current_kv
|
||||
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
|
||||
if fixed_kv:
|
||||
next_key_values[i].advance(seq_len)
|
||||
|
||||
# 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 +921,9 @@ class Llama2_(nn.Module):
|
|||
if i == intermediate_output:
|
||||
intermediate = x.clone()
|
||||
|
||||
if prefetch_queue is not None:
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
|
|
@ -810,7 +937,7 @@ class Llama2_(nn.Module):
|
|||
if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
|
||||
intermediate = self.norm(intermediate)
|
||||
|
||||
if len(next_key_values) > 0:
|
||||
if next_key_values:
|
||||
return x, intermediate, next_key_values
|
||||
else:
|
||||
return x, intermediate
|
||||
|
|
@ -874,12 +1001,7 @@ class BaseGenerate:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
import torch
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
import comfy.ops
|
||||
from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR
|
||||
from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt
|
||||
|
||||
|
||||
MODEL_CONFIG = {
|
||||
"vocab_size": 200000,
|
||||
"hidden_size": 4096,
|
||||
"intermediate_size": 12288,
|
||||
"num_hidden_layers": 36,
|
||||
"num_attention_heads": 32,
|
||||
"num_key_value_heads": 8,
|
||||
"max_position_embeddings": 10240,
|
||||
"rms_norm_eps": 1e-6,
|
||||
"rope_theta": 1000000.0,
|
||||
"head_dim": 128,
|
||||
"audio_vocab_size": 1024,
|
||||
"audio_num_codebooks": 8,
|
||||
"decoder_num_heads": 16,
|
||||
"decoder_intermediate_size": 6144,
|
||||
"decoder_num_layers": 4,
|
||||
}
|
||||
|
||||
|
||||
def detect_merged_config(state_dict, prefix=""):
|
||||
return {
|
||||
"merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
|
||||
"merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
|
||||
"decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
|
||||
"decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
|
||||
}
|
||||
|
||||
|
||||
class MiniMaxMusic3Tokenizer:
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
tokenizer_json = tokenizer_data.get("tokenizer_json")
|
||||
if tokenizer_json is None:
|
||||
raise ValueError("MiniMax Music3 text encoder checkpoint is missing tokenizer_json")
|
||||
if torch.is_tensor(tokenizer_json):
|
||||
tokenizer_json = tokenizer_json.detach().cpu().numpy().tobytes()
|
||||
self.tokenizer_json = tokenizer_json
|
||||
self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8"))
|
||||
for token, expected in SPECIAL_TOKEN_IDS.items():
|
||||
if self.tokenizer.token_to_id(token) != expected:
|
||||
raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}")
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
|
||||
prompt = build_prompt(text, kwargs.get("lyrics", ""))
|
||||
token_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids
|
||||
return {
|
||||
"minimax_music3": [[(token, 1.0) for token in token_ids]],
|
||||
"seed": int(kwargs.get("seed", 0)),
|
||||
"max_audio_frames": int(kwargs.get("max_audio_frames", MAX_AUDIO_FRAMES)),
|
||||
"cfg_scale": float(kwargs.get("cfg_scale", CFG_SCALE)),
|
||||
"top_k": int(kwargs.get("top_k", CFG_TOP_K)),
|
||||
}
|
||||
|
||||
def state_dict(self):
|
||||
return {"tokenizer_json": torch.frombuffer(bytearray(self.tokenizer_json), dtype=torch.uint8)}
|
||||
|
||||
def decode(self, token_ids, skip_special_tokens=True):
|
||||
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
|
||||
class MiniMaxMusic3TEModel(MiniMaxMusic3AR):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}, projection_config=None):
|
||||
dtype = torch.bfloat16
|
||||
quant_config = model_options.get("quantization_metadata", None)
|
||||
operations = model_options.get("custom_operations", None)
|
||||
if operations is None:
|
||||
operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast
|
||||
super().__init__({**MODEL_CONFIG, **(projection_config or {})}, dtype, device, operations)
|
||||
self.dtypes = {dtype}
|
||||
self.execution_device = device
|
||||
|
||||
def set_clip_options(self, options):
|
||||
self.execution_device = options.get("execution_device", self.execution_device)
|
||||
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
def get_dynamic_vram__units(self):
|
||||
units, last_units = self.model.get_dynamic_vram__units()
|
||||
if self.model.pruned_embedding:
|
||||
last_units = [*last_units, self.model.embed_tokens_prefill]
|
||||
return [(self.model.audio_decoder, self.model.audio_extra_embedding), *units], last_units
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
token_ids = [token for token, _ in token_weight_pairs["minimax_music3"][0]]
|
||||
input_ids = torch.tensor([token_ids], dtype=torch.long)
|
||||
seed = token_weight_pairs["seed"]
|
||||
max_audio_frames = token_weight_pairs["max_audio_frames"]
|
||||
cfg_scale = token_weight_pairs["cfg_scale"]
|
||||
top_k = token_weight_pairs["top_k"]
|
||||
hidden = self.generate(input_ids, seed, max_audio_frames, self.execution_device, cfg_scale, top_k)
|
||||
return hidden.unsqueeze(0), None, {}
|
||||
|
||||
def load_state_dict(self, state_dict, strict=True, assign=False):
|
||||
if self.model.pruned_embedding is None:
|
||||
self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict
|
||||
if self.model.pruned_embedding:
|
||||
del self.model.embed_tokens
|
||||
else:
|
||||
del self.model.embed_tokens_prefill, self.model.embed_tokens_audio
|
||||
if self.model.pruned_lm_head is None:
|
||||
self.model.pruned_lm_head = "model.lm_head_pruned.weight" in state_dict
|
||||
if self.model.pruned_lm_head:
|
||||
del self.model.lm_head
|
||||
else:
|
||||
del self.model.lm_head_pruned
|
||||
return super().load_state_dict(state_dict, strict=strict, assign=assign)
|
||||
|
||||
def load_sd(self, state_dict):
|
||||
return self.load_state_dict(state_dict, strict=False, assign=getattr(self, "can_assign_sd", False))
|
||||
|
|
@ -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(...)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -6,7 +6,13 @@ from typing_extensions import override
|
|||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.bria import (
|
||||
BriaEditImageRequest,
|
||||
BriaEraseRequest,
|
||||
BriaExpandRequest,
|
||||
BriaExpandResponse,
|
||||
BriaGenFillRequest,
|
||||
BriaImageEditResponse,
|
||||
BriaImageResultResponse,
|
||||
BriaIncreaseResolutionRequest,
|
||||
BriaRemoveBackgroundRequest,
|
||||
BriaRemoveBackgroundResponse,
|
||||
BriaRemoveVideoBackgroundRequest,
|
||||
|
|
@ -21,13 +27,30 @@ from comfy_api_nodes.util import (
|
|||
convert_mask_to_image,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor_by_max_side,
|
||||
get_image_dimensions,
|
||||
poll_op,
|
||||
sync_op,
|
||||
upload_image_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_string,
|
||||
validate_video_duration,
|
||||
)
|
||||
|
||||
BRIA_MAX_OUTPUT_SIDE = 8192
|
||||
BRIA_MIN_RATIO = 0.5
|
||||
BRIA_MAX_RATIO = 3.0
|
||||
BRIA_MIN_SHORT_SIDE = 224
|
||||
|
||||
|
||||
def _upscaled_output_side(height: int, width: int, multiplier: int) -> int:
|
||||
prescale = max(1.0, BRIA_MIN_SHORT_SIDE / min(height, width))
|
||||
return round(max(height, width) * prescale * multiplier)
|
||||
|
||||
|
||||
def _smallest_output_side(height: int, width: int, multiplier: int) -> int:
|
||||
return round(max(height, width) / min(height, width) * BRIA_MIN_SHORT_SIDE * multiplier)
|
||||
|
||||
|
||||
class BriaImageEditNode(IO.ComfyNode):
|
||||
|
||||
|
|
@ -243,6 +266,503 @@ class BriaRemoveImageBackground(IO.ComfyNode):
|
|||
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
|
||||
|
||||
|
||||
def _mask_to_binary_image(mask: Input.Image, action: str) -> torch.Tensor:
|
||||
binary = (mask > 0.5).float()
|
||||
if not binary.any():
|
||||
raise ValueError(
|
||||
f"The mask is empty, so there is nothing to {action}. Masks are binarized at 50%: "
|
||||
f"areas painted at less than half opacity are ignored."
|
||||
)
|
||||
return convert_mask_to_image(binary)
|
||||
|
||||
|
||||
def _validate_mask_aspect_ratio(image: Input.Image, mask: Input.Image) -> None:
|
||||
ih, iw = image.shape[1], image.shape[2]
|
||||
mh, mw = mask.shape[-2], mask.shape[-1]
|
||||
if abs(iw * mh - ih * mw) > 0.01 * ih * mw:
|
||||
raise ValueError(f"Mask must have the same aspect ratio as the image: image is {iw}x{ih}, mask is {mw}x{mh}.")
|
||||
|
||||
|
||||
class BriaGenFill(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="BriaGenFill",
|
||||
display_name="Bria Generative Fill",
|
||||
category="partner/image/Bria",
|
||||
description="Generate objects or scenery inside a masked region of an image using Bria.",
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.Mask.Input(
|
||||
"mask",
|
||||
tooltip="White areas are filled with generated content, black areas are preserved. "
|
||||
"The mask is binarized before sending, so partially painted areas count as white. "
|
||||
"Must have the same aspect ratio as the image.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of what to generate inside the masked region.",
|
||||
),
|
||||
IO.String.Input("negative_prompt", multiline=True, default=""),
|
||||
IO.Boolean.Input(
|
||||
"refine_prompt",
|
||||
default=True,
|
||||
tooltip="Automatically adjust the prompt for better results; "
|
||||
"disable to use the prompt exactly as written.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=1,
|
||||
max=2147483647,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"moderation",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("false", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"true",
|
||||
[
|
||||
IO.Boolean.Input("prompt_content_moderation", default=False),
|
||||
IO.Boolean.Input("visual_input_moderation", default=False),
|
||||
IO.Boolean.Input("visual_output_moderation", default=False),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Moderation settings",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Image.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.0429}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
mask: Input.Image,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
refine_prompt: bool,
|
||||
seed: int,
|
||||
moderation: InputModerationSettings,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, min_length=1)
|
||||
_validate_mask_aspect_ratio(image, mask)
|
||||
mask_image = _mask_to_binary_image(mask, "fill")
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/bria/v2/image/edit/gen_fill", method="POST"),
|
||||
data=BriaGenFillRequest(
|
||||
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
|
||||
mask=await upload_image_to_comfyapi(
|
||||
cls, mask_image, total_pixels=None, wait_label="Uploading mask"
|
||||
),
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
refine_prompt=refine_prompt,
|
||||
seed=seed,
|
||||
prompt_content_moderation=moderation.get("prompt_content_moderation", False),
|
||||
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
|
||||
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
|
||||
),
|
||||
response_model=BriaStatusResponse,
|
||||
)
|
||||
response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
|
||||
status_extractor=lambda r: r.status,
|
||||
response_model=BriaImageResultResponse,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
|
||||
|
||||
|
||||
class BriaEraser(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="BriaEraser",
|
||||
display_name="Bria Eraser",
|
||||
category="partner/image/Bria",
|
||||
description="Remove objects or areas outlined by a mask from an image using Bria.",
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.Mask.Input(
|
||||
"mask",
|
||||
tooltip="White areas are erased, black areas are preserved. "
|
||||
"The mask is binarized before sending, so partially painted areas count as white. "
|
||||
"Must have the same aspect ratio as the image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"mask_type",
|
||||
options=["manual", "automatic"],
|
||||
tooltip="manual for hand-drawn or brush masks, "
|
||||
"automatic for masks produced by segmentation models such as SAM.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"moderation",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("false", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"true",
|
||||
[
|
||||
IO.Boolean.Input("visual_input_moderation", default=False),
|
||||
IO.Boolean.Input("visual_output_moderation", default=False),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Moderation settings",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Image.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.0286}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
mask: Input.Image,
|
||||
mask_type: str,
|
||||
moderation: dict,
|
||||
) -> IO.NodeOutput:
|
||||
_validate_mask_aspect_ratio(image, mask)
|
||||
mask_image = _mask_to_binary_image(mask, "erase")
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/bria/v2/image/edit/erase", method="POST"),
|
||||
data=BriaEraseRequest(
|
||||
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
|
||||
mask=await upload_image_to_comfyapi(
|
||||
cls, mask_image, total_pixels=None, wait_label="Uploading mask"
|
||||
),
|
||||
mask_type=mask_type,
|
||||
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
|
||||
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
|
||||
),
|
||||
response_model=BriaStatusResponse,
|
||||
)
|
||||
response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
|
||||
status_extractor=lambda r: r.status,
|
||||
response_model=BriaImageResultResponse,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
|
||||
|
||||
|
||||
class BriaExpandImage(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="BriaExpandImage",
|
||||
display_name="Bria Expand Image",
|
||||
category="partner/image/Bria",
|
||||
description="Expand an image beyond its borders with generated content using Bria.",
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.DynamicCombo.Input(
|
||||
"expand_mode",
|
||||
options=[
|
||||
*[IO.DynamicCombo.Option(ratio, []) for ratio in
|
||||
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"]],
|
||||
IO.DynamicCombo.Option(
|
||||
"custom_ratio",
|
||||
[
|
||||
IO.Int.Input(
|
||||
"ratio_width",
|
||||
default=21,
|
||||
min=1,
|
||||
max=100,
|
||||
tooltip="Width side of the target ratio: 21 and 9 give 21:9.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"ratio_height",
|
||||
default=9,
|
||||
min=1,
|
||||
max=100,
|
||||
tooltip="Height side of the target ratio: 21 and 9 give 21:9. "
|
||||
f"Bria only accepts width/height between {BRIA_MIN_RATIO} and "
|
||||
f"{BRIA_MAX_RATIO}, so anything taller than 1:2 needs the manual mode.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"manual",
|
||||
[
|
||||
IO.Int.Input("canvas_width", default=1000, min=64, max=5000),
|
||||
IO.Int.Input("canvas_height", default=1000, min=64, max=5000),
|
||||
IO.Int.Input(
|
||||
"image_width",
|
||||
default=500,
|
||||
min=1,
|
||||
max=5000,
|
||||
tooltip="Width of the original image inside the canvas.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"image_height",
|
||||
default=500,
|
||||
min=1,
|
||||
max=5000,
|
||||
tooltip="Height of the original image inside the canvas.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"image_x",
|
||||
default=250,
|
||||
min=-5000,
|
||||
max=5000,
|
||||
tooltip="X position of the image's top-left corner inside the canvas; "
|
||||
"may fall outside the canvas, cropping the image.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"image_y",
|
||||
default=250,
|
||||
min=-5000,
|
||||
max=5000,
|
||||
tooltip="Y position of the image's top-left corner inside the canvas; "
|
||||
"may fall outside the canvas, cropping the image.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Target shape of the expanded image: a preset aspect ratio, a custom ratio, "
|
||||
"or manual placement of the original image on a canvas. "
|
||||
"Manual is the only mode that can reach a canvas taller than 1:2.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional description of the expanded scene; "
|
||||
"when empty, Bria generates one from the image.",
|
||||
),
|
||||
IO.String.Input("negative_prompt", multiline=True, default=""),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=1,
|
||||
max=2147483647,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"moderation",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("false", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"true",
|
||||
[
|
||||
IO.Boolean.Input("prompt_content_moderation", default=False),
|
||||
IO.Boolean.Input("visual_input_moderation", default=False),
|
||||
IO.Boolean.Input("visual_output_moderation", default=False),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Moderation settings",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
IO.String.Output(display_name="prompt", tooltip="The prompt used for the expansion; "
|
||||
"auto-generated by Bria when the prompt input is empty."),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.0286}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
expand_mode: dict,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
seed: int,
|
||||
moderation: InputModerationSettings,
|
||||
) -> IO.NodeOutput:
|
||||
mode = expand_mode["expand_mode"]
|
||||
aspect_ratio = canvas_size = original_image_size = original_image_location = None
|
||||
if mode == "manual":
|
||||
canvas_size = [expand_mode["canvas_width"], expand_mode["canvas_height"]]
|
||||
original_image_size = [expand_mode["image_width"], expand_mode["image_height"]]
|
||||
original_image_location = [expand_mode["image_x"], expand_mode["image_y"]]
|
||||
elif mode == "custom_ratio":
|
||||
ratio_width, ratio_height = expand_mode["ratio_width"], expand_mode["ratio_height"]
|
||||
aspect_ratio = ratio_width / ratio_height
|
||||
if not BRIA_MIN_RATIO <= aspect_ratio <= BRIA_MAX_RATIO:
|
||||
raise ValueError(
|
||||
f"Bria accepts a width-to-height ratio between {BRIA_MIN_RATIO} and {BRIA_MAX_RATIO}: "
|
||||
f"{ratio_width}:{ratio_height} is {aspect_ratio:.4f}. "
|
||||
f"Use the manual expand mode to reach a canvas of any shape."
|
||||
)
|
||||
else:
|
||||
aspect_ratio = mode
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/bria/v2/image/edit/expand", method="POST"),
|
||||
data=BriaExpandRequest(
|
||||
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
|
||||
aspect_ratio=aspect_ratio,
|
||||
canvas_size=canvas_size,
|
||||
original_image_size=original_image_size,
|
||||
original_image_location=original_image_location,
|
||||
prompt=prompt if prompt else None,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
seed=seed,
|
||||
prompt_content_moderation=moderation.get("prompt_content_moderation", False),
|
||||
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
|
||||
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
|
||||
),
|
||||
response_model=BriaStatusResponse,
|
||||
)
|
||||
response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
|
||||
status_extractor=lambda r: r.status,
|
||||
response_model=BriaExpandResponse,
|
||||
)
|
||||
return IO.NodeOutput(
|
||||
await download_url_to_image_tensor(response.result.image_url),
|
||||
response.result.prompt or "",
|
||||
)
|
||||
|
||||
|
||||
class BriaIncreaseResolution(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="BriaIncreaseResolution",
|
||||
display_name="Bria Increase Resolution",
|
||||
category="partner/image/Bria",
|
||||
description="Upscale an image by 2x or 4x using Bria, preserving the original content.",
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.Combo.Input(
|
||||
"desired_increase",
|
||||
options=["2", "4"],
|
||||
tooltip="Resolution multiplier. The output must fit within 8192 pixels on each side.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"auto_downscale",
|
||||
default=False,
|
||||
tooltip="Automatically lower the multiplier, and downscale the input image if that is "
|
||||
"still not enough, when the output would exceed the limit.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"moderation",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("false", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"true",
|
||||
[
|
||||
IO.Boolean.Input("visual_input_moderation", default=False),
|
||||
IO.Boolean.Input("visual_output_moderation", default=False),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Moderation settings",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Image.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.0286}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
desired_increase: str,
|
||||
auto_downscale: bool,
|
||||
moderation: dict,
|
||||
) -> IO.NodeOutput:
|
||||
multiplier = int(desired_increase)
|
||||
height, width = get_image_dimensions(image)
|
||||
if _upscaled_output_side(height, width, multiplier) > BRIA_MAX_OUTPUT_SIDE:
|
||||
candidates = [c for c in (4, 2) if c <= multiplier]
|
||||
if not auto_downscale:
|
||||
predicted = _upscaled_output_side(height, width, multiplier)
|
||||
raise ValueError(
|
||||
f"Bria can upscale up to a maximum output dimension of {BRIA_MAX_OUTPUT_SIDE} pixels: "
|
||||
f"input is {width}x{height}, x{multiplier} would be {predicted} pixels on the long side. "
|
||||
f"Enable auto_downscale, or use a smaller input image or a lower multiplier."
|
||||
)
|
||||
fitted = next(
|
||||
(c for c in candidates if _upscaled_output_side(height, width, c) <= BRIA_MAX_OUTPUT_SIDE), None
|
||||
)
|
||||
if fitted is not None:
|
||||
multiplier = fitted
|
||||
else:
|
||||
shrinkable = next((c for c in sorted(candidates) if _smallest_output_side(height, width, c)
|
||||
<= BRIA_MAX_OUTPUT_SIDE), None)
|
||||
if shrinkable is None:
|
||||
raise ValueError(
|
||||
f"This image cannot be upscaled by Bria at any multiplier: it is {width}x{height}, and "
|
||||
f"Bria first enlarges the short side to {BRIA_MIN_SHORT_SIDE} pixels, which pushes the "
|
||||
f"long side past the {BRIA_MAX_OUTPUT_SIDE} pixel limit. Crop it to a squarer shape first."
|
||||
)
|
||||
multiplier = shrinkable
|
||||
image = downscale_image_tensor_by_max_side(image, max_side=BRIA_MAX_OUTPUT_SIDE // multiplier)
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/bria/v2/image/edit/increase_resolution", method="POST"),
|
||||
data=BriaIncreaseResolutionRequest(
|
||||
image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
|
||||
desired_increase=multiplier,
|
||||
visual_input_content_moderation=moderation.get("visual_input_moderation", False),
|
||||
visual_output_content_moderation=moderation.get("visual_output_moderation", False),
|
||||
),
|
||||
response_model=BriaStatusResponse,
|
||||
)
|
||||
response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
|
||||
status_extractor=lambda r: r.status,
|
||||
response_model=BriaImageResultResponse,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
|
||||
|
||||
|
||||
class BriaRemoveVideoBackground(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
|
|
@ -572,6 +1092,10 @@ class BriaExtension(ComfyExtension):
|
|||
return [
|
||||
BriaImageEditNode,
|
||||
BriaRemoveImageBackground,
|
||||
BriaGenFill,
|
||||
BriaEraser,
|
||||
BriaExpandImage,
|
||||
BriaIncreaseResolution,
|
||||
BriaRemoveVideoBackground,
|
||||
BriaVideoGreenScreen,
|
||||
BriaVideoReplaceBackground,
|
||||
|
|
|
|||
|
|
@ -3,12 +3,14 @@ from typing import Optional
|
|||
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 @@ HAILUO_03_QUERY_ENDPOINT = "/proxy/minimax/v2/query/video_generation" # + /{tas
|
|||
HAILUO_03_MODELS = {"MiniMax H3": "MiniMax-H3"}
|
||||
HAILUO_03_FAILED_STATUSES = ["failed", "cancelled", "expired"]
|
||||
|
||||
HAILUO_03_CONTEXT_IR_ENDPOINT = "/proxy/minimax/v2/h3_context_ir"
|
||||
HAILUO_03_REGENERATION_ENDPOINT = "/proxy/minimax/v2/video_regeneration"
|
||||
|
||||
|
||||
def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = True):
|
||||
inputs = [
|
||||
|
|
@ -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 @@ class MinimaxHailuo03ReferenceNode(IO.ComfyNode):
|
|||
)
|
||||
|
||||
|
||||
class MinimaxHailuo03ContextIRNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="MinimaxHailuo03ContextIRNode",
|
||||
display_name="MiniMax H3 Context IR (Prompt Enhancer)",
|
||||
category="partner/video/MiniMax",
|
||||
description="Analyze text and media context with MiniMax H3 Context IR and produce an enhanced, "
|
||||
"structured video prompt. Feed the output into the prompt of a MiniMax H3 video node and attach "
|
||||
"the same media there in the same order, because the enhanced prompt refers to the attached "
|
||||
"media by position.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"MiniMax H3",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of the video you intend to generate.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=5,
|
||||
min=4,
|
||||
max=15,
|
||||
step=1,
|
||||
tooltip="Duration of the video you intend to generate, in seconds (4-15).",
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"ratio",
|
||||
options=["adaptive", "16:9", "4:3", "1:1", "3:4", "9:16", "21:9"],
|
||||
default="adaptive",
|
||||
tooltip="Aspect ratio of the video you intend to generate. 'adaptive' "
|
||||
"requires at least one image, video, or audio input.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("reference_image"),
|
||||
names=[
|
||||
"image_1",
|
||||
"image_2",
|
||||
"image_3",
|
||||
"image_4",
|
||||
"image_5",
|
||||
"image_6",
|
||||
"image_7",
|
||||
"image_8",
|
||||
"image_9",
|
||||
],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Subject or style reference images, referred to in the prompt "
|
||||
"as 'Image 1'..'Image 9' in connection order. Up to 9 images.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_videos",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Video.Input("reference_video"),
|
||||
names=["video_1", "video_2", "video_3"],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Motion or scene reference videos, referred to in the prompt "
|
||||
"as 'Video 1'..'Video 3' in connection order. Up to 3 videos, "
|
||||
"2-15 seconds each, 15 seconds in total.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_audios",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Audio.Input("reference_audio"),
|
||||
names=["audio_1", "audio_2", "audio_3"],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Audio references, referred to in the prompt as "
|
||||
"'Audio 1'..'Audio 3' in connection order. Up to 3 clips, "
|
||||
"2-15 seconds each, 15 seconds in total. Cannot be used without "
|
||||
"a reference image or video.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="Model to use for prompt enhancement.",
|
||||
),
|
||||
IO.Image.Input(
|
||||
"first_frame",
|
||||
tooltip="First frame of the video you intend to generate. Cannot be combined with "
|
||||
"reference media.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Image.Input(
|
||||
"last_frame",
|
||||
tooltip="Last frame of the video you intend to generate. Cannot be combined with "
|
||||
"reference media.",
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.String.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(
|
||||
inputs=["first_frame", "last_frame"],
|
||||
input_groups=["model.reference_images", "model.reference_videos", "model.reference_audios"],
|
||||
),
|
||||
expr="""
|
||||
(
|
||||
$imgsRaw := $lookup(inputGroups, "model.reference_images");
|
||||
$imgs := $imgsRaw ? $imgsRaw : 0;
|
||||
$vidsRaw := $lookup(inputGroups, "model.reference_videos");
|
||||
$vids := $vidsRaw ? $vidsRaw : 0;
|
||||
$audsRaw := $lookup(inputGroups, "model.reference_audios");
|
||||
$auds := $audsRaw ? $audsRaw : 0;
|
||||
$frames := (inputs.first_frame.connected ? 1 : 0) + (inputs.last_frame.connected ? 1 : 0);
|
||||
($imgs + $vids + $auds) > 0
|
||||
? {"type": "range_usd", "min_usd": 0.06, "max_usd": 0.11, "format": {"approximate": true}}
|
||||
: $frames > 0
|
||||
? {"type": "usd", "usd": 0.05, "format": {"approximate": true}}
|
||||
: {"type": "usd", "usd": 0.02, "format": {"approximate": true}}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
model: dict,
|
||||
first_frame: torch.Tensor | None = None,
|
||||
last_frame: torch.Tensor | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(model["prompt"], strip_whitespace=True, min_length=1)
|
||||
|
||||
reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
|
||||
reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
|
||||
reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
|
||||
has_frames = first_frame is not None or last_frame is not None
|
||||
has_references = bool(reference_images) or bool(reference_videos) or bool(reference_audios)
|
||||
if has_frames and has_references:
|
||||
raise ValueError(
|
||||
"First/last frame and reference media are mutually exclusive. Use frames for an "
|
||||
"image-to-video prompt, or reference media for a reference-to-video prompt."
|
||||
)
|
||||
if reference_audios and not reference_images and not reference_videos:
|
||||
raise ValueError("Reference audio cannot be used without a reference image or video.")
|
||||
if not has_frames and not has_references and model["ratio"] == "adaptive":
|
||||
raise ValueError(
|
||||
"Ratio 'adaptive' is not supported for text-only requests; select an explicit aspect ratio."
|
||||
)
|
||||
|
||||
for frame in (first_frame, last_frame):
|
||||
if frame is not None:
|
||||
validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
|
||||
validate_image_dimensions(frame, min_width=256, min_height=256)
|
||||
for image in reference_images.values():
|
||||
validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
|
||||
validate_image_dimensions(image, min_width=256, min_height=256)
|
||||
|
||||
total_video_duration = 0.0
|
||||
for i, video in enumerate(reference_videos.values(), 1):
|
||||
try:
|
||||
fps = float(video.get_frame_rate())
|
||||
except Exception:
|
||||
fps = 0.0
|
||||
if fps and not (23.9 <= fps <= 60.5):
|
||||
raise ValueError(f"Reference video {i} is {fps:.2f} FPS. Supported range is 23.976-60 FPS.")
|
||||
try:
|
||||
dur = video.get_duration()
|
||||
except Exception:
|
||||
continue
|
||||
if dur < 1.8:
|
||||
raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
|
||||
total_video_duration += dur
|
||||
if total_video_duration > 15.1:
|
||||
raise ValueError(
|
||||
f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
|
||||
)
|
||||
|
||||
total_audio_duration = 0.0
|
||||
for i, audio in enumerate(reference_audios.values(), 1):
|
||||
dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
|
||||
if dur < 1.8:
|
||||
raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
|
||||
total_audio_duration += dur
|
||||
if total_audio_duration > 15.1:
|
||||
raise ValueError(
|
||||
f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
|
||||
)
|
||||
|
||||
content: list = [Hailuo03TextContent(text=model["prompt"])]
|
||||
if first_frame is not None:
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, first_frame, max_images=1, wait_label="Uploading first frame"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="first_frame",
|
||||
)
|
||||
)
|
||||
if last_frame is not None:
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, last_frame, max_images=1, wait_label="Uploading last frame"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="last_frame",
|
||||
)
|
||||
)
|
||||
for i, image in enumerate(reference_images.values(), 1):
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, image, max_images=1, wait_label=f"Uploading image {i}"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="reference_image",
|
||||
)
|
||||
)
|
||||
for i, video in enumerate(reference_videos.values(), 1):
|
||||
content.append(
|
||||
Hailuo03VideoContent(
|
||||
video_url=Hailuo03VideoContentUrl(
|
||||
url=await upload_video_to_comfyapi(cls, video, wait_label=f"Uploading video {i}"),
|
||||
),
|
||||
)
|
||||
)
|
||||
for audio in reference_audios.values():
|
||||
content.append(
|
||||
Hailuo03AudioContent(
|
||||
audio_url=Hailuo03AudioContentUrl(
|
||||
url=await upload_audio_to_comfyapi(
|
||||
cls,
|
||||
audio,
|
||||
container_format="mp3",
|
||||
codec_name="libmp3lame",
|
||||
mime_type="audio/mpeg",
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=HAILUO_03_CONTEXT_IR_ENDPOINT, method="POST"),
|
||||
response_model=Hailuo03TaskCreationResponse,
|
||||
data=Hailuo03ContextIRRequest(
|
||||
model=HAILUO_03_MODELS[model["model"]],
|
||||
content=content,
|
||||
duration=model["duration"],
|
||||
ratio=None if model["ratio"] == "adaptive" else model["ratio"],
|
||||
),
|
||||
)
|
||||
task_result = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
|
||||
response_model=Hailuo03TaskQueryResponse,
|
||||
status_extractor=lambda r: r.task.status,
|
||||
failed_statuses=HAILUO_03_FAILED_STATUSES,
|
||||
poll_interval=5,
|
||||
)
|
||||
prompt = task_result.task.content.prompt if task_result.task.content else None
|
||||
if not prompt:
|
||||
raise Exception(f"No enhanced prompt in the response: {task_result.model_dump()}")
|
||||
return IO.NodeOutput(prompt)
|
||||
|
||||
|
||||
class MinimaxHailuo03RegenerateNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="MinimaxHailuo03RegenerateNode",
|
||||
display_name="MiniMax H3 Regenerate to 2K",
|
||||
category="partner/video/MiniMax",
|
||||
description="Re-render a MiniMax H3 768P output at 2K resolution. Connect the unmodified 768P "
|
||||
"video and the exact prompt used to generate it; if the original generation used first/last "
|
||||
"frames or reference media, attach the same inputs.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"MiniMax H3",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="The exact prompt used to generate the source video.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["2K"],
|
||||
tooltip="Resolution to re-render the source video at.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("reference_image"),
|
||||
names=[
|
||||
"image_1",
|
||||
"image_2",
|
||||
"image_3",
|
||||
"image_4",
|
||||
"image_5",
|
||||
"image_6",
|
||||
"image_7",
|
||||
"image_8",
|
||||
"image_9",
|
||||
],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Reference images from the original generation, in the same "
|
||||
"order. Up to 9 images.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_videos",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Video.Input("reference_video"),
|
||||
names=["video_1", "video_2", "video_3"],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Reference videos from the original generation, in the same "
|
||||
"order. Up to 3 videos, 2-15 seconds each, 15 seconds in total.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_audios",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Audio.Input("reference_audio"),
|
||||
names=["audio_1", "audio_2", "audio_3"],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Audio references from the original generation, in the same "
|
||||
"order. Up to 3 clips, 2-15 seconds each, 15 seconds in total. "
|
||||
"Cannot be used without a reference image or video.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="Model to use for video regeneration.",
|
||||
),
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="The MiniMax H3 768P output video to re-render. Connect the unmodified output "
|
||||
"of a MiniMax H3 video node (24 FPS, 4-15 seconds). 2K outputs cannot be used.",
|
||||
),
|
||||
IO.Image.Input(
|
||||
"first_frame",
|
||||
tooltip="First frame image from the original generation, if one was used.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Image.Input(
|
||||
"last_frame",
|
||||
tooltip="Last frame image from the original generation, if one was used.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"watermark",
|
||||
default=False,
|
||||
tooltip="Whether to add an AIGC watermark to the video.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type": "usd", "usd": 0.0715, "format": {"suffix": "/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
model: dict,
|
||||
video: Input.Video,
|
||||
watermark: bool,
|
||||
first_frame: torch.Tensor | None = None,
|
||||
last_frame: torch.Tensor | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(model["prompt"], strip_whitespace=True, min_length=1)
|
||||
|
||||
try:
|
||||
fps = float(video.get_frame_rate())
|
||||
except Exception:
|
||||
fps = 0.0
|
||||
if fps and not (23.9 <= fps <= 24.1):
|
||||
raise ValueError(
|
||||
f"The source video is {fps:.2f} FPS. Regeneration accepts unmodified MiniMax H3 768P "
|
||||
"outputs, which are 24 FPS."
|
||||
)
|
||||
try:
|
||||
width, height = video.get_dimensions()
|
||||
except Exception:
|
||||
width = height = 0
|
||||
if width and height and (width % 32 or height % 32 or width * height > 1_032_192):
|
||||
raise ValueError(
|
||||
f"The source video is {width}x{height}. Regeneration accepts MiniMax H3 768P outputs "
|
||||
"(width and height divisible by 32, at most 1,032,192 total pixels); 2K outputs cannot "
|
||||
"be used as a source."
|
||||
)
|
||||
try:
|
||||
frame_count = video.get_frame_count()
|
||||
except Exception:
|
||||
frame_count = 0
|
||||
if frame_count and (frame_count < 107 or frame_count > 362 or (frame_count - 107) % 17):
|
||||
raise ValueError(
|
||||
f"The source video has {frame_count} frames. Regeneration accepts unmodified "
|
||||
"MiniMax H3 outputs, whose length is 107 to 362 frames in steps of 17 "
|
||||
"(4 to 15 seconds at 24 FPS)."
|
||||
)
|
||||
|
||||
reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
|
||||
reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
|
||||
reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
|
||||
if (first_frame is not None or last_frame is not None) and (
|
||||
reference_images or reference_videos or reference_audios
|
||||
):
|
||||
raise ValueError(
|
||||
"First/last frame and reference media are mutually exclusive. Use frames for an "
|
||||
"image-to-video prompt, or reference media for a reference-to-video prompt."
|
||||
)
|
||||
if reference_audios and not reference_images and not reference_videos:
|
||||
raise ValueError("Reference audio cannot be used without a reference image or video.")
|
||||
|
||||
for frame in (first_frame, last_frame):
|
||||
if frame is not None:
|
||||
validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
|
||||
validate_image_dimensions(frame, min_width=256, min_height=256)
|
||||
for image in reference_images.values():
|
||||
validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False) # 0.4 to 2.5
|
||||
validate_image_dimensions(image, min_width=256, min_height=256)
|
||||
|
||||
total_video_duration = 0.0
|
||||
for i, ref_video in enumerate(reference_videos.values(), 1):
|
||||
try:
|
||||
ref_fps = float(ref_video.get_frame_rate())
|
||||
except Exception:
|
||||
ref_fps = 0.0
|
||||
if ref_fps and not (23.9 <= ref_fps <= 60.5):
|
||||
raise ValueError(f"Reference video {i} is {ref_fps:.2f} FPS. Supported range is 23.976-60 FPS.")
|
||||
try:
|
||||
dur = ref_video.get_duration()
|
||||
except Exception:
|
||||
continue
|
||||
if dur < 1.8:
|
||||
raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
|
||||
total_video_duration += dur
|
||||
if total_video_duration > 15.1:
|
||||
raise ValueError(
|
||||
f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
|
||||
)
|
||||
|
||||
total_audio_duration = 0.0
|
||||
for i, audio in enumerate(reference_audios.values(), 1):
|
||||
dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
|
||||
if dur < 1.8:
|
||||
raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
|
||||
total_audio_duration += dur
|
||||
if total_audio_duration > 15.1:
|
||||
raise ValueError(
|
||||
f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
|
||||
)
|
||||
|
||||
content: list = [
|
||||
Hailuo03VideoContent(
|
||||
video_url=Hailuo03VideoContentUrl(
|
||||
url=await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video"),
|
||||
),
|
||||
role="base_video",
|
||||
),
|
||||
Hailuo03TextContent(text=model["prompt"]),
|
||||
]
|
||||
if first_frame is not None:
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, first_frame, max_images=1, wait_label="Uploading first frame"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="first_frame",
|
||||
)
|
||||
)
|
||||
if last_frame is not None:
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, last_frame, max_images=1, wait_label="Uploading last frame"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="last_frame",
|
||||
)
|
||||
)
|
||||
for i, image in enumerate(reference_images.values(), 1):
|
||||
content.append(
|
||||
Hailuo03ImageContent(
|
||||
image_url=Hailuo03ImageContentUrl(
|
||||
url=(
|
||||
await upload_images_to_comfyapi(
|
||||
cls, image, max_images=1, wait_label=f"Uploading image {i}"
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
role="reference_image",
|
||||
)
|
||||
)
|
||||
for i, ref_video in enumerate(reference_videos.values(), 1):
|
||||
content.append(
|
||||
Hailuo03VideoContent(
|
||||
video_url=Hailuo03VideoContentUrl(
|
||||
url=await upload_video_to_comfyapi(cls, ref_video, wait_label=f"Uploading video {i}"),
|
||||
),
|
||||
)
|
||||
)
|
||||
for audio in reference_audios.values():
|
||||
content.append(
|
||||
Hailuo03AudioContent(
|
||||
audio_url=Hailuo03AudioContentUrl(
|
||||
url=await upload_audio_to_comfyapi(
|
||||
cls,
|
||||
audio,
|
||||
container_format="mp3",
|
||||
codec_name="libmp3lame",
|
||||
mime_type="audio/mpeg",
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=HAILUO_03_REGENERATION_ENDPOINT, method="POST"),
|
||||
response_model=Hailuo03TaskCreationResponse,
|
||||
data=Hailuo03RegenerationRequest(
|
||||
model=HAILUO_03_MODELS[model["model"]],
|
||||
content=content,
|
||||
resolution=model["resolution"],
|
||||
aigc_watermark=watermark,
|
||||
),
|
||||
)
|
||||
task_result = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
|
||||
response_model=Hailuo03TaskQueryResponse,
|
||||
status_extractor=lambda r: r.task.status,
|
||||
failed_statuses=HAILUO_03_FAILED_STATUSES,
|
||||
poll_interval=10,
|
||||
)
|
||||
video_url = task_result.task.content.url if task_result.task.content else None
|
||||
if not video_url:
|
||||
raise Exception(f"No video URL in the response: {task_result.model_dump()}")
|
||||
return IO.NodeOutput(await download_url_to_video_output(video_url))
|
||||
|
||||
|
||||
class MinimaxExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
|
|
@ -950,6 +1541,8 @@ class MinimaxExtension(ComfyExtension):
|
|||
MinimaxHailuo03TextToVideoNode,
|
||||
MinimaxHailuo03FirstLastFrameNode,
|
||||
MinimaxHailuo03ReferenceNode,
|
||||
MinimaxHailuo03ContextIRNode,
|
||||
MinimaxHailuo03RegenerateNode,
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import comfy.model_sampling
|
|||
import comfy.nested_tensor
|
||||
import comfy.utils
|
||||
import node_helpers
|
||||
from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
CANVAS_MULTIPLE = 32
|
||||
|
|
@ -67,6 +68,16 @@ def _resize(image, width, height, crop):
|
|||
return samples.movedim(1, -1)
|
||||
|
||||
|
||||
def _encode_ref_audio(audio_vae, audio):
|
||||
waveform = audio["waveform"] # [B, C, L]
|
||||
sr = audio["sample_rate"]
|
||||
vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
|
||||
if sr != vae_sr:
|
||||
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
|
||||
z = audio_vae.encode(waveform[:1].movedim(1, -1)) # [1, 32, 2, T]
|
||||
return z, z.shape[-1]
|
||||
|
||||
|
||||
def _empty_av_latent(width, height, length, batch_size=1):
|
||||
frame_count, latent_t, audio_t = temporal_shape(length)
|
||||
video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
|
||||
|
|
@ -144,13 +155,87 @@ class MiniMaxH3ImageToVideo(io.ComfyNode):
|
|||
if keyframes:
|
||||
for kf in keyframes:
|
||||
kf["latent"] = vae.encode(kf.pop("image"))
|
||||
cond = node_helpers.conditioning_set_values(cond, {
|
||||
"minimax_keyframes": keyframes,
|
||||
"minimax_frame_count": frame_count,
|
||||
})
|
||||
cond = node_helpers.conditioning_set_values(cond, {"minimax_keyframes": keyframes})
|
||||
return io.NodeOutput(cond, latent)
|
||||
|
||||
|
||||
class MiniMaxH3AddGuide(io.ComfyNode):
|
||||
"""Anchor image and/or audio guides at an arbitrary pixel frame of the target video."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="MiniMaxH3AddGuide",
|
||||
display_name="Add Guide for MiniMax H3",
|
||||
category="model/conditioning/minimax",
|
||||
description="Anchor an image, a short clip, audio, or a clip with its soundtrack at any frame of a MiniMax H3 video. Chain several nodes to anchor several frames.",
|
||||
inputs=[
|
||||
io.Conditioning.Input("positive"),
|
||||
io.Vae.Input("vae", optional=True, tooltip="Video VAE, needed when an image is connected."),
|
||||
io.Vae.Input("audio_vae", optional=True, tooltip="Audio VAE, needed when an audio is connected."),
|
||||
io.Latent.Input("latent"),
|
||||
io.Image.Input("image", optional=True, tooltip="Image or video frames to anchor. Multi-frame batches are anchored as a clip and cropped down to the model's valid clip lengths: 5, 22, 39... (17k + 5) frames. Batches shorter than 5 frames use only the first image."),
|
||||
io.Audio.Input("audio", optional=True,
|
||||
tooltip="Soundtrack to anchor starting at the same frame index, cropped to the video's remaining duration."),
|
||||
io.Int.Input("frame_idx", default=0, min=-9999, max=9999,
|
||||
tooltip="Frame index to anchor the image or the clip's first frame at. Negative values are counted from the end of the video."),
|
||||
],
|
||||
outputs=[io.Conditioning.Output(display_name="positive")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, positive, latent, frame_idx, vae=None, audio_vae=None, image=None, audio=None) -> io.NodeOutput:
|
||||
samples = latent["samples"]
|
||||
if not samples.is_nested or len(samples.tensors) != 2 or samples.tensors[0].ndim != 5 or samples.tensors[0].shape[1] != 24:
|
||||
raise ValueError("MiniMaxH3AddGuide expects a MiniMax H3 AV latent")
|
||||
if image is None and audio is None:
|
||||
raise ValueError("MiniMaxH3AddGuide needs an image or an audio to anchor")
|
||||
video = samples.tensors[0]
|
||||
height = video.shape[3] * 16
|
||||
width = video.shape[4] * 16
|
||||
frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(video.shape[2]))
|
||||
|
||||
guide_frames = 1
|
||||
if image is not None:
|
||||
if vae is None:
|
||||
raise ValueError("anchoring guide frames needs the vae input")
|
||||
guide_frames = image.shape[0]
|
||||
if guide_frames < 5:
|
||||
guide_frames = 1
|
||||
else:
|
||||
while guide_frames % 17 != 5:
|
||||
guide_frames -= 1
|
||||
|
||||
resolved_frame_index = frame_idx if frame_idx >= 0 else frame_count + frame_idx
|
||||
if resolved_frame_index < 0 or resolved_frame_index + guide_frames > frame_count:
|
||||
if guide_frames == 1:
|
||||
raise ValueError("frame_idx {} is outside the video's {} frames".format(frame_idx, frame_count))
|
||||
raise ValueError("a {} frame guide clip at frame_idx {} does not fit in the video's {} frames".format(
|
||||
guide_frames, frame_idx, frame_count))
|
||||
|
||||
keyframe = {"resolved_frame_index": resolved_frame_index}
|
||||
if image is not None:
|
||||
frames = _resize(image[:guide_frames], width, height, "center")
|
||||
keyframe["latent"] = vae.encode(frames)
|
||||
|
||||
if audio is not None:
|
||||
if audio_vae is None:
|
||||
raise ValueError("anchoring guide audio needs the audio_vae input")
|
||||
audio_latent, audio_rt = _encode_ref_audio(audio_vae, audio)
|
||||
# the streams share one time axis: FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
|
||||
max_rt = math.floor(samples.tensors[1].shape[-1] - FRAME_RESCALE * resolved_frame_index)
|
||||
if max_rt < 1:
|
||||
raise ValueError("frame_idx {} is past the end of the video's audio track".format(frame_idx))
|
||||
if audio_rt > max_rt:
|
||||
audio_latent = audio_latent[..., :max_rt].clone()
|
||||
keyframe["audio_latent"] = audio_latent
|
||||
|
||||
keyframes = list(positive[0][1].get("minimax_keyframes", []))
|
||||
keyframes.append(keyframe)
|
||||
positive = node_helpers.conditioning_set_values(positive, {"minimax_keyframes": keyframes})
|
||||
return io.NodeOutput(positive)
|
||||
|
||||
|
||||
class MiniMaxH3ReferenceToVideo(io.ComfyNode):
|
||||
"""ref2va: prompt + reference images / videos / audio -> conditioning + AV latent.
|
||||
|
||||
|
|
@ -197,16 +282,6 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
|
|||
outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _encode_ref_audio(audio_vae, audio):
|
||||
waveform = audio["waveform"] # [B, C, L]
|
||||
sr = audio["sample_rate"]
|
||||
vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
|
||||
if sr != vae_sr:
|
||||
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
|
||||
z = audio_vae.encode(waveform[:1].movedim(1, -1)) # [1, 32, 2, T]
|
||||
return z, z.shape[-1]
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_size="match",
|
||||
ref_images=None, ref_videos=None, ref_video_audios=None, ref_audios=None) -> io.NodeOutput:
|
||||
|
|
@ -254,7 +329,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
|
|||
z = vae.encode(frames)
|
||||
audio_latent, ref_audio_t = (None, 0)
|
||||
if soundtrack is not None:
|
||||
audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, soundtrack)
|
||||
audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, soundtrack)
|
||||
# the soundtrack gets its own <Audio j> label, emitted before <Video k>
|
||||
ref_items.append({"type": "audio"})
|
||||
# Qwen sees the video at 2 fps with timestamps
|
||||
|
|
@ -269,7 +344,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
|
|||
for audio in (ref_audios or {}).values():
|
||||
if audio is None:
|
||||
continue
|
||||
audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, audio)
|
||||
audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, audio)
|
||||
ref_items.append({"type": "audio"})
|
||||
ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
|
||||
|
||||
|
|
@ -329,6 +404,7 @@ class MiniMaxH3Extension(ComfyExtension):
|
|||
return [
|
||||
EmptyMiniMaxH3LatentAV,
|
||||
MiniMaxH3ImageToVideo,
|
||||
MiniMaxH3AddGuide,
|
||||
MiniMaxH3ReferenceToVideo,
|
||||
MiniMaxH3SigmaShift,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -29,7 +29,7 @@ class PreviewAny():
|
|||
value = str(source)
|
||||
elif source is not None:
|
||||
try:
|
||||
value = json.dumps(source, indent=4)
|
||||
value = json.dumps(source, indent=4, ensure_ascii=False)
|
||||
except Exception:
|
||||
try:
|
||||
value = str(source)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
# This file is automatically generated by the build process when version is
|
||||
# updated in pyproject.toml.
|
||||
__version__ = "0.32.0"
|
||||
__version__ = "0.33.0"
|
||||
|
|
|
|||
2
main.py
2
main.py
|
|
@ -248,7 +248,7 @@ import hook_breaker_ac10a0
|
|||
import comfy.memory_management
|
||||
import comfy.model_patcher
|
||||
|
||||
if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()):
|
||||
if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia()):
|
||||
if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
|
||||
logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
|
||||
else:
|
||||
|
|
|
|||
6
nodes.py
6
nodes.py
|
|
@ -290,6 +290,9 @@ class ConditioningZeroOut:
|
|||
conditioning_lyrics = d.get("conditioning_lyrics", None)
|
||||
if conditioning_lyrics is not None:
|
||||
d["conditioning_lyrics"] = torch.zeros_like(conditioning_lyrics)
|
||||
conditioning_scale = d.get("conditioning_scale", None)
|
||||
if conditioning_scale is not None:
|
||||
d["conditioning_scale"] = torch.zeros_like(conditioning_scale)
|
||||
n = [torch.zeros_like(t[0]), d]
|
||||
c.append(n)
|
||||
return (c, )
|
||||
|
|
@ -1015,7 +1018,7 @@ class CLIPLoader:
|
|||
|
||||
CATEGORY = "model/loaders"
|
||||
|
||||
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
||||
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",
|
||||
|
|
|
|||
321
openapi.yaml
321
openapi.yaml
|
|
@ -35,6 +35,10 @@ components:
|
|||
description: Timestamp when the asset was last accessed
|
||||
format: date-time
|
||||
type: string
|
||||
loader_path:
|
||||
description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
|
||||
nullable: true
|
||||
type: string
|
||||
metadata:
|
||||
additionalProperties: true
|
||||
description: System-managed metadata from download sources (HuggingFace, CivitAI, etc.) - read-only, not user-modifiable
|
||||
|
|
@ -165,6 +169,10 @@ components:
|
|||
format: uuid
|
||||
nullable: true
|
||||
type: string
|
||||
loader_path:
|
||||
description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
|
||||
nullable: true
|
||||
type: string
|
||||
mime_type:
|
||||
description: Updated MIME type of the asset
|
||||
type: string
|
||||
|
|
@ -188,6 +196,31 @@ components:
|
|||
- id
|
||||
- updated_at
|
||||
type: object
|
||||
ChurnkeyAuthResponse:
|
||||
description: |
|
||||
Credentials the Churnkey embed requires to launch the cancel flow.
|
||||
`auth_hash` is hex-encoded HMAC-SHA256 of `customer_id` signed with the
|
||||
server's CHURNKEY_HMAC_SECRET; it is bound to that single customer ID
|
||||
and must not be reused for other customers.
|
||||
properties:
|
||||
auth_hash:
|
||||
description: Hex-encoded HMAC-SHA256(customer_id, CHURNKEY_HMAC_SECRET)
|
||||
type: string
|
||||
customer_id:
|
||||
description: Stripe customer ID for the workspace
|
||||
type: string
|
||||
mode:
|
||||
description: Churnkey environment matching the configured app
|
||||
enum:
|
||||
- live
|
||||
- test
|
||||
- sandbox
|
||||
type: string
|
||||
required:
|
||||
- customer_id
|
||||
- auth_hash
|
||||
- mode
|
||||
type: object
|
||||
CreateWorkflowRequest:
|
||||
description: Request body for creating a new saved workflow.
|
||||
properties:
|
||||
|
|
@ -511,6 +544,25 @@ components:
|
|||
required:
|
||||
- history
|
||||
type: object
|
||||
JobAssetsResponse:
|
||||
description: Paginated list of the assets produced by a single job.
|
||||
properties:
|
||||
assets:
|
||||
description: The job's output assets for the requested page (empty when the job produced none)
|
||||
items:
|
||||
$ref: '#/components/schemas/JobOutputAsset'
|
||||
type: array
|
||||
job_id:
|
||||
description: ID of the job these assets belong to
|
||||
format: uuid
|
||||
type: string
|
||||
pagination:
|
||||
$ref: '#/components/schemas/PaginationInfo'
|
||||
required:
|
||||
- job_id
|
||||
- assets
|
||||
- pagination
|
||||
type: object
|
||||
JobCancelResponse:
|
||||
description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
|
||||
properties:
|
||||
|
|
@ -565,6 +617,9 @@ components:
|
|||
additionalProperties: true
|
||||
description: Primary preview output (only for terminal states)
|
||||
type: object
|
||||
previewable_outputs_count:
|
||||
description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
|
||||
type: integer
|
||||
status:
|
||||
description: User-friendly job status
|
||||
enum:
|
||||
|
|
@ -597,6 +652,13 @@ components:
|
|||
workflow_id:
|
||||
description: UUID identifying the workflow graph definition
|
||||
type: string
|
||||
workflow_version_id:
|
||||
description: |
|
||||
UUID of the cloud workflow version this job is pinned to, if the
|
||||
submission carried one (see PromptRequest's workflow_version_id).
|
||||
Absent for jobs submitted without that association, including
|
||||
every job submitted through the public API v2 today.
|
||||
type: string
|
||||
workspace_id:
|
||||
description: |
|
||||
ID of the workspace that owns this job. A successful (200)
|
||||
|
|
@ -645,6 +707,9 @@ components:
|
|||
additionalProperties: true
|
||||
description: Primary preview output (only present for terminal states)
|
||||
type: object
|
||||
previewable_outputs_count:
|
||||
description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
|
||||
type: integer
|
||||
status:
|
||||
description: User-friendly job status
|
||||
enum:
|
||||
|
|
@ -662,6 +727,56 @@ components:
|
|||
- status
|
||||
- create_time
|
||||
type: object
|
||||
JobOutputAsset:
|
||||
description: |
|
||||
An asset produced by a job, enriched with the per-output node context
|
||||
(`node_id`, `output_key`, `output_index`) correlated from the job's
|
||||
execution outputs by content hash. The node-context fields are null
|
||||
when the asset cannot be matched to an output entry.
|
||||
properties:
|
||||
created_at:
|
||||
description: Timestamp when the asset was created
|
||||
format: date-time
|
||||
type: string
|
||||
hash:
|
||||
description: Blake3 hash of the asset content.
|
||||
pattern: ^blake3:[a-f0-9]{64}$
|
||||
type: string
|
||||
id:
|
||||
description: Unique identifier for the asset
|
||||
format: uuid
|
||||
type: string
|
||||
mime_type:
|
||||
description: MIME type of the asset
|
||||
type: string
|
||||
name:
|
||||
description: Name of the asset file
|
||||
type: string
|
||||
node_id:
|
||||
description: ID of the workflow node that produced this asset, if known
|
||||
nullable: true
|
||||
type: string
|
||||
output_index:
|
||||
description: Zero-based index of this asset within the node's output slot, if known
|
||||
nullable: true
|
||||
type: integer
|
||||
output_key:
|
||||
description: Output slot key under the producing node (e.g. "images"), if known
|
||||
nullable: true
|
||||
type: string
|
||||
preview_url:
|
||||
description: Relative URL for asset preview/thumbnail
|
||||
format: uri-reference
|
||||
type: string
|
||||
size:
|
||||
description: Size of the asset in bytes
|
||||
format: int64
|
||||
type: integer
|
||||
required:
|
||||
- id
|
||||
- name
|
||||
- created_at
|
||||
type: object
|
||||
JobStatusResponse:
|
||||
description: Job status information
|
||||
properties:
|
||||
|
|
@ -1522,7 +1637,11 @@ paths:
|
|||
operationId: listAssets
|
||||
parameters:
|
||||
- deprecated: true
|
||||
description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
|
||||
description: |
|
||||
Deprecated alias for `tags_all`, kept permanently for existing
|
||||
callers. Filter assets that have ALL of these tags. Combining it
|
||||
with `tags_all`, or exceeding 100 tags (counted after removing
|
||||
empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
|
||||
explode: false
|
||||
in: query
|
||||
name: include_tags
|
||||
|
|
@ -1532,7 +1651,11 @@ paths:
|
|||
type: array
|
||||
style: form
|
||||
- deprecated: true
|
||||
description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
|
||||
description: |
|
||||
Deprecated alias for `tags_none`, kept permanently for existing
|
||||
callers. Exclude assets that have ANY of these tags. Combining it
|
||||
with `tags_none`, or exceeding 100 tags (counted after removing
|
||||
empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
|
||||
explode: false
|
||||
in: query
|
||||
name: exclude_tags
|
||||
|
|
@ -1541,7 +1664,14 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Filter assets that have ALL of these tags
|
||||
- description: |
|
||||
Filter assets that have ALL of these tags. Tag values are opaque
|
||||
byte-strings compared exactly and case-sensitively; unknown tags
|
||||
are not an error — they simply match nothing. Replaces the
|
||||
deprecated `include_tags`. Sending both spellings, listing the
|
||||
same tag here and in `tags_none`, or exceeding 100 tags per list
|
||||
(counted after removing empty values and duplicates) returns 400
|
||||
`INVALID_TAG_FILTER`.
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_all
|
||||
|
|
@ -1550,7 +1680,14 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Filter assets that have AT LEAST ONE of these tags
|
||||
- description: |
|
||||
Filter assets that have AT LEAST ONE of these tags. Combines with
|
||||
`tags_all`/`tags_none` by intersection (`tags_none` always wins;
|
||||
overlap with `tags_none` is allowed and leaves a dead term).
|
||||
Supplying a positive tag filter (`tags_any`, `tags_all`, or
|
||||
`include_tags`) replaces the default category filter that is
|
||||
otherwise applied. Lists over 100 tags (counted after removing
|
||||
empty values and duplicates) return 400 `INVALID_TAG_FILTER`.
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_any
|
||||
|
|
@ -1559,7 +1696,11 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Exclude assets that have ANY of these tags
|
||||
- description: |
|
||||
Exclude assets that have ANY of these tags. Replaces the
|
||||
deprecated `exclude_tags`. Sending both spellings, or exceeding
|
||||
100 tags per list (counted after removing empty values and
|
||||
duplicates), returns 400 `INVALID_TAG_FILTER`.
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_none
|
||||
|
|
@ -2342,7 +2483,11 @@ paths:
|
|||
operationId: getAssetTagHistogram
|
||||
parameters:
|
||||
- deprecated: true
|
||||
description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
|
||||
description: |
|
||||
Deprecated alias for `tags_all`, kept permanently for existing
|
||||
callers. Filter assets that have ALL of these tags. The same
|
||||
combination and list-size rules as on `/api/assets` apply
|
||||
(400 `INVALID_TAG_FILTER`).
|
||||
explode: false
|
||||
in: query
|
||||
name: include_tags
|
||||
|
|
@ -2352,7 +2497,11 @@ paths:
|
|||
type: array
|
||||
style: form
|
||||
- deprecated: true
|
||||
description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
|
||||
description: |
|
||||
Deprecated alias for `tags_none`, kept permanently for existing
|
||||
callers. Exclude assets that have ANY of these tags. The same
|
||||
combination and list-size rules as on `/api/assets` apply
|
||||
(400 `INVALID_TAG_FILTER`).
|
||||
explode: false
|
||||
in: query
|
||||
name: exclude_tags
|
||||
|
|
@ -2361,7 +2510,10 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Filter assets that have ALL of these tags
|
||||
- description: |
|
||||
Filter assets that have ALL of these tags. Replaces the deprecated
|
||||
`include_tags`. The same combination and list-size rules as on
|
||||
`/api/assets` apply (400 `INVALID_TAG_FILTER`).
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_all
|
||||
|
|
@ -2370,7 +2522,11 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Filter assets that have AT LEAST ONE of these tags
|
||||
- description: |
|
||||
Filter assets that have AT LEAST ONE of these tags. Combines with
|
||||
`tags_all`/`tags_none` by intersection (`tags_none` always wins).
|
||||
The same combination and list-size rules as on `/api/assets` apply
|
||||
(400 `INVALID_TAG_FILTER`).
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_any
|
||||
|
|
@ -2379,7 +2535,10 @@ paths:
|
|||
type: string
|
||||
type: array
|
||||
style: form
|
||||
- description: Exclude assets that have ANY of these tags
|
||||
- description: |
|
||||
Exclude assets that have ANY of these tags. Replaces the deprecated
|
||||
`exclude_tags`. The same combination and list-size rules as on
|
||||
`/api/assets` apply (400 `INVALID_TAG_FILTER`).
|
||||
explode: false
|
||||
in: query
|
||||
name: tags_none
|
||||
|
|
@ -2440,6 +2599,49 @@ paths:
|
|||
summary: Get tag histogram for filtered assets
|
||||
tags:
|
||||
- file
|
||||
/api/billing/churnkey/auth:
|
||||
get:
|
||||
description: |
|
||||
Returns the Stripe customer identifier and a server-signed
|
||||
HMAC-SHA256 of the customer ID, used to launch the Churnkey-hosted
|
||||
cancellation flow embed.
|
||||
operationId: getChurnkeyAuth
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ChurnkeyAuthResponse'
|
||||
description: Success
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Workspace has no Stripe customer (never subscribed)
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
"503":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Churnkey is not configured on the server
|
||||
security:
|
||||
- BearerAuth: []
|
||||
summary: Get Churnkey HMAC auth credentials
|
||||
tags:
|
||||
- billing
|
||||
/api/embeddings:
|
||||
get:
|
||||
description: Returns the list of text-encoder embeddings available on disk.
|
||||
|
|
@ -2460,9 +2662,10 @@ paths:
|
|||
Returns a list of model folders available in the system.
|
||||
This is an experimental endpoint that replaces the legacy /models endpoint.
|
||||
Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
|
||||
Once the model_type migration is active the names are model_type folder_names
|
||||
(e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
|
||||
its directory path.
|
||||
The folder vocabulary is resolved per request from the caller's identity: where the
|
||||
model_type migration is active for that caller the names are model_type folder_names
|
||||
(e.g. `ultralytics_bbox`), and a folder with no folder_name mapping is returned by its
|
||||
directory path. An authenticated response can therefore differ from an anonymous one.
|
||||
operationId: getModelFolders
|
||||
responses:
|
||||
"200":
|
||||
|
|
@ -2479,7 +2682,10 @@ paths:
|
|||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
security: []
|
||||
security:
|
||||
- ApiKeyAuth: []
|
||||
- BearerAuth: []
|
||||
- {}
|
||||
summary: Get available model folders
|
||||
tags:
|
||||
- file
|
||||
|
|
@ -2488,6 +2694,10 @@ paths:
|
|||
description: |
|
||||
Returns a list of models available in the specified folder.
|
||||
This is an experimental endpoint that provides enhanced model information.
|
||||
Accepted folder identifiers are those returned by /api/experiment/models for the same
|
||||
caller. That vocabulary is request-scoped, so list folders and fetch a folder's models
|
||||
with the same credentials — a name obtained anonymously may not resolve when
|
||||
authenticated, and vice versa.
|
||||
operationId: getModelsInFolder
|
||||
parameters:
|
||||
- description: The folder name to list models from
|
||||
|
|
@ -2518,7 +2728,10 @@ paths:
|
|||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
security: []
|
||||
security:
|
||||
- ApiKeyAuth: []
|
||||
- BearerAuth: []
|
||||
- {}
|
||||
summary: Get models in a specific folder
|
||||
tags:
|
||||
- file
|
||||
|
|
@ -3155,6 +3368,74 @@ paths:
|
|||
summary: Get full job details
|
||||
tags:
|
||||
- workflow
|
||||
/api/jobs/{job_id}/assets:
|
||||
get:
|
||||
description: |
|
||||
Retrieve a paginated list of the assets produced by a specific job,
|
||||
enriched with the per-output node context (`node_id`, `output_key`,
|
||||
`output_index`) correlated from the job's execution outputs by content
|
||||
hash. Unlike `GET /api/assets?job_ids={id}`, this endpoint is scoped to a
|
||||
single job and carries node-level placement, making it suited to job
|
||||
output views rather than the general asset browser. Returns an empty
|
||||
`assets` array for jobs that produced no assets.
|
||||
operationId: getJobAssets
|
||||
parameters:
|
||||
- description: Job identifier (UUID)
|
||||
in: path
|
||||
name: job_id
|
||||
required: true
|
||||
schema:
|
||||
format: uuid
|
||||
type: string
|
||||
- description: Maximum number of assets to return (1-500)
|
||||
in: query
|
||||
name: limit
|
||||
schema:
|
||||
default: 20
|
||||
maximum: 500
|
||||
minimum: 1
|
||||
type: integer
|
||||
- description: Number of assets to skip for pagination
|
||||
in: query
|
||||
name: offset
|
||||
schema:
|
||||
default: 0
|
||||
minimum: 0
|
||||
type: integer
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/JobAssetsResponse'
|
||||
description: Success - Job assets returned
|
||||
"400":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Invalid request parameters
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized - Authentication required
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Job not found or does not belong to the user
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
summary: List a job's output assets
|
||||
tags:
|
||||
- workflow
|
||||
/api/jobs/{job_id}/cancel:
|
||||
post:
|
||||
description: |
|
||||
|
|
@ -3360,6 +3641,12 @@ paths:
|
|||
schema:
|
||||
$ref: '#/components/schemas/PromptErrorResponse'
|
||||
description: Payment required - Insufficient credits
|
||||
"403":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/PromptErrorResponse'
|
||||
description: Workspace governance policy blocks one or more partner providers (error.type PARTNER_NODE_DISABLED; error.class_types lists the offending nodes, error.providers the disabled providers)
|
||||
"413":
|
||||
content:
|
||||
application/json:
|
||||
|
|
@ -3371,7 +3658,7 @@ paths:
|
|||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/PromptErrorResponse'
|
||||
description: Payment required - User has not paid
|
||||
description: 'Retryable backpressure. Two distinct causes, disambiguated by the body''s `error.type`, NOT by parsing `error.message`: `PAYMENT_REQUIRED` / `FREE_TIER_UNAVAILABLE` / `FREE_TIER_EXHAUSTED` / `PARTNER_NODE_PAYMENT_REQUIRED` (a billing gate - retrying without paying never succeeds), or `QUEUE_LIMIT` (this workspace''s bounded job queue is full - retrying after some queued jobs complete will succeed).'
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
|
|
@ -5210,6 +5497,8 @@ tags:
|
|||
name: user
|
||||
- description: Background task management
|
||||
name: task
|
||||
- description: Workspace billing and subscription management
|
||||
name: billing
|
||||
- description: Workflow storage and version management
|
||||
name: workflows
|
||||
- description: Job queue state and control
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "ComfyUI"
|
||||
version = "0.32.0"
|
||||
version = "0.33.0"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
comfyui-frontend-package==1.48.7
|
||||
comfyui-workflow-templates==0.11.39
|
||||
comfyui-workflow-templates==0.11.41
|
||||
comfyui-embedded-docs==0.5.9
|
||||
torch
|
||||
torchsde
|
||||
|
|
@ -22,7 +22,7 @@ alembic
|
|||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.30
|
||||
comfy-kitchen==0.2.31
|
||||
comfy-aimdo==0.4.13
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,30 @@
|
|||
from unittest.mock import patch, MagicMock
|
||||
|
||||
mock_nodes = MagicMock()
|
||||
mock_nodes.MAX_RESOLUTION = 16384
|
||||
mock_server = MagicMock()
|
||||
|
||||
with patch.dict("sys.modules", {"nodes": mock_nodes, "server": mock_server}):
|
||||
from comfy_extras.nodes_preview_any import PreviewAny
|
||||
|
||||
|
||||
class TestPreviewAnyMain:
|
||||
@staticmethod
|
||||
def _exec(source) -> dict:
|
||||
return PreviewAny().main(source)
|
||||
|
||||
def test_dict_keeps_non_ascii(self):
|
||||
result = self._exec({"greeting": "你好"})
|
||||
assert "你好" in result["ui"]["text"][0]
|
||||
assert "\\u" not in result["ui"]["text"][0]
|
||||
assert result["result"][0] == result["ui"]["text"][0]
|
||||
|
||||
def test_list_keeps_non_ascii(self):
|
||||
result = self._exec(["你好", "こんにちは"])
|
||||
assert "こんにちは" in result["result"][0]
|
||||
assert "\\u" not in result["result"][0]
|
||||
|
||||
def test_string_passthrough(self):
|
||||
result = self._exec("你好")
|
||||
assert result["ui"]["text"][0] == "你好"
|
||||
assert result["result"][0] == "你好"
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
"""Gemma4 chat template regression tests."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
args.cpu = True
|
||||
|
||||
import comfy.text_encoders.gemma4 as gemma4 # noqa: E402
|
||||
|
||||
PROMPT = "describe a cute anime girl with fennec ears"
|
||||
THOUGHT_BLOCK = "<|channel>thought\n<channel|>"
|
||||
|
||||
# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a
|
||||
# closed thought block when thinking is off.
|
||||
NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B]
|
||||
PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B]
|
||||
|
||||
|
||||
class _CaptureTemplate:
|
||||
"""Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files."""
|
||||
llama_text = ""
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
|
||||
self.llama_text = text
|
||||
return {}
|
||||
|
||||
|
||||
def build_template(variant, **kwargs):
|
||||
prime = variant.tokenizer.tokenizer_class.prime_empty_thought
|
||||
probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})()
|
||||
probe.tokenize_with_weights(PROMPT, **kwargs)
|
||||
return probe.llama_text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
|
||||
def test_thinking_enabled_only_asks_via_the_system_turn(variant):
|
||||
template = build_template(variant, skip_template=False, thinking=True)
|
||||
assert template == f"<|turn>system\n<|think|>\n<turn|>\n<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", NO_PRIMING)
|
||||
def test_thinking_disabled_does_not_prime_a_thought_channel(variant):
|
||||
template = build_template(variant, skip_template=False, thinking=False)
|
||||
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
|
||||
assert "channel" not in template
|
||||
assert "<|think|>" not in template
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", PRIMING)
|
||||
def test_thinking_disabled_primes_a_thought_channel(variant):
|
||||
template = build_template(variant, skip_template=False, thinking=False)
|
||||
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n{THOUGHT_BLOCK}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
|
||||
@pytest.mark.parametrize("thinking", [False, True])
|
||||
def test_skip_template_passes_text_through_unchanged(variant, thinking):
|
||||
assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT
|
||||
Loading…
Reference in New Issue