diff --git a/AGENTS.md b/AGENTS.md index 05efd834b..20014ce7e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,9 @@ better to remove a broken feature path than keep a complicated partial fix. - Preserve existing APIs, node names, model-loading behavior, file layout, and workflow compatibility unless the change is explicitly about replacing them. +- When compatibility is explicitly out of scope, remove compatibility-only + aliases, duplicate nodes, legacy entry points, and preset wrappers instead of + retaining parallel ways to perform the same operation. - Code must look hand-written for this repository. Changes that read like generic AI-generated code will be rejected automatically: unnecessary helper layers, vague names, boilerplate comments, defensive branches without a real @@ -96,6 +99,13 @@ unless they are read by current code and change current behavior. Remove pass-through or stored-but-unused values instead of preserving upstream or deprecated API baggage. +- Do not add a model-specific option to a shared helper when only one caller + needs it. Keep one-off behavior at the model integration boundary, or extend + the shared helper only when the option is a coherent reusable capability. +- Implementations of shared model interfaces should accept the standard caller + contract without model-specific rejection branches for optional capabilities + they do not consume. Let supported behavior be determined by implementation + paths that actually use those inputs. - If an implementation needs auxiliary values for its own workflow, expose them through a private helper or a clearly named implementation-specific method instead of overloading the public method's return contract. @@ -154,6 +164,10 @@ `comfy-kitchen` helpers where they already solve the problem. - Use optimized comfy-kitchen ops in places where they improve performance without changing the expected dtype, device, memory, or interface behavior. +- Prefer ComfyUI's shared optimized kernels and backend dispatchers over + handwritten implementations of the same operation. Remove duplicate local + kernels and adapt inputs to the shared operation's documented layout while + preserving the model's original math and output contract. - All models should use the optimized attention function selected by ComfyUI. Treat optimized backend functions, dispatch helpers, and capability-selected callables as opaque. Higher-level code must not inspect function identity, @@ -176,6 +190,12 @@ - Model detection code that inspects linear weight shapes should only use the first dimension. The second dimension may be half the original size for NVFP4 or other 4-bit quantized models. +- A model-detection signature must guard every state-dict key it dereferences. + Do not partially match a format and then raise an incidental `KeyError` while + extracting its configuration. +- Order model-detection checks from established or more-specific signatures to + newer or broader signatures. Put a broad new detector near the generic + fallback when giving it higher precedence could steal another model family. - Avoid adding `einops` usage in core inference code. Use native torch tensor ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`, `unsqueeze`, and `squeeze` instead. @@ -192,11 +212,23 @@ methods for scalar or structural calculations. - Avoid unnecessary casts and transfers. Preserve the intended compute dtype, storage dtype, bias dtype, and original tensor shape metadata. +- Do not cast the result of an optimized backend operation back to its input + dtype unless that backend's documented result contract requires normalization. + In particular, trust the selected optimized-attention implementation to honor + its dtype contract. - Keep model-native latent layout handling inside the model or latent-format owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent dimensions in nodes or other caller-side adapters just to satisfy a model forward; the model path should consume and return the native latent shape for that model family. +- DiT models should accept latent dimensions that are not exact patch-size + multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified + target or reference input, then crop only the target output back to its + original dimensions. +- Avoid defensive shape and configuration checks that merely replace the clear + failure from the tensor operation immediately below them. Add explicit + validation only when it provides materially better context at a real boundary + or prevents silent incorrect output. - Assume inputs to the main model forward are already in the compute dtype by default, except integer inputs such as some model timestep tensors. Do not add defensive or convenience casts in model code; it is better for invalid dtype @@ -260,6 +292,15 @@ - Model implementations should add the minimal number of ComfyUI nodes required to run the model. Reuse existing nodes as much as possible; adapting the model to work with existing nodes is strongly preferred over creating new nodes. +- Use `io.Autogrow` for a variable number of repeated inputs instead of a fixed + series of numbered optional sockets. Set its minimum to zero when the model + has a valid no-item path, and cap it only when the model has a real limit. +- Mark inputs optional when execution has a valid path that does not read them. + If one optional input is needed only to process another optional input, do not + force users on the path that supplies neither to connect it. +- Conditioning nodes should normally output conditioning only. Do not expose + input or intermediate images as convenience outputs for downstream sizing or + routing; use the existing image path or a dedicated image operation instead. - Nodes should output only values they own. Do not add pass-through outputs for workflow convenience unless the node is explicitly an output node. Existing models, latents, conditioning, or other inputs should flow directly to the diff --git a/comfy/ldm/joyimage/model.py b/comfy/ldm/joyimage/model.py new file mode 100644 index 000000000..bca12c391 --- /dev/null +++ b/comfy/ldm/joyimage/model.py @@ -0,0 +1,445 @@ +# https://github.com/jdopensource/JoyAI-Image-Edit (Apache 2.0) +import math +from typing import Optional, Tuple + +import comfy_kitchen +import torch +import torch.nn as nn + +import comfy.ldm.common_dit +import comfy.ops +import comfy.patcher_extension +from comfy.ldm.lightricks.model import GELU_approx, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from comfy.ldm.modules.attention import optimized_attention + + +class JoyImageModulate(nn.Module): + def __init__(self, hidden_size: int, factor: int, dtype=None, device=None): + super().__init__() + self.factor = factor + self.modulate_table = nn.Parameter( + torch.empty(1, factor, hidden_size, dtype=dtype, device=device) + ) + + def forward(self, x: torch.Tensor) -> list: + if x.ndim != 3: + x = x.unsqueeze(1) + table = comfy.ops.cast_to_input(self.modulate_table, x) + return [o.squeeze(1) for o in (table + x).chunk(self.factor, dim=1)] + + +class JoyImageFeedForward(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + self.net = nn.ModuleList([ + GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations), + nn.Identity(), + operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device), + ]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for module in self.net: + x = module(x) + return x + + +class JoyImageAttention(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + inner_dim = num_attention_heads * attention_head_dim + + self.img_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device) + self.img_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device) + self.img_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device) + self.img_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device) + + self.txt_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device) + self.txt_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device) + self.txt_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device) + self.txt_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device) + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + image_rotary_emb: torch.Tensor, + transformer_options=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + heads = self.num_attention_heads + + img_q, img_k, img_v = self.img_attn_qkv(img).chunk(3, dim=-1) + txt_q, txt_k, txt_v = self.txt_attn_qkv(txt).chunk(3, dim=-1) + + img_q = img_q.unflatten(-1, (heads, -1)) + img_k = img_k.unflatten(-1, (heads, -1)) + img_v = img_v.unflatten(-1, (heads, -1)) + txt_q = txt_q.unflatten(-1, (heads, -1)) + txt_k = txt_k.unflatten(-1, (heads, -1)) + txt_v = txt_v.unflatten(-1, (heads, -1)) + + img_q = self.img_attn_q_norm(img_q) + img_k = self.img_attn_k_norm(img_k) + txt_q = self.txt_attn_q_norm(txt_q) + txt_k = self.txt_attn_k_norm(txt_k) + + img_q, img_k = comfy_kitchen.apply_rope(img_q, img_k, image_rotary_emb) + + joint_q = torch.cat([img_q, txt_q], dim=1) + joint_k = torch.cat([img_k, txt_k], dim=1) + joint_v = torch.cat([img_v, txt_v], dim=1) + + joint_q = joint_q.flatten(2, 3) + joint_k = joint_k.flatten(2, 3) + joint_v = joint_v.flatten(2, 3) + + joint_out = optimized_attention(joint_q, joint_k, joint_v, heads=heads, transformer_options=transformer_options) + + seq_img = img.shape[1] + img_out = joint_out[:, :seq_img, :] + txt_out = joint_out[:, seq_img:, :] + + img_out = self.img_attn_proj(img_out) + txt_out = self.txt_attn_proj(txt_out) + return img_out, txt_out + + +class JoyImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: float = 4.0, + eps: float = 1e-6, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + mlp_hidden_dim = int(dim * mlp_width_ratio) + + self.img_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device) + self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device) + self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device) + self.img_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations) + + self.txt_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device) + self.txt_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device) + self.txt_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device) + self.txt_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations) + + self.attn = JoyImageAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + eps=eps, + dtype=dtype, + device=device, + operations=operations, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: torch.Tensor, + transformer_options=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(temb) + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(temb) + + img_normed = self.img_norm1(hidden_states) + txt_normed = self.txt_norm1(encoder_hidden_states) + img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1) + txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1) + + img_attn, txt_attn = self.attn(img_modulated, txt_modulated, image_rotary_emb, transformer_options=transformer_options) + + hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1) + + img_ffn_normed = self.img_norm2(hidden_states) + txt_ffn_normed = self.txt_norm2(encoder_hidden_states) + img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1) + txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1) + hidden_states = hidden_states + self.img_mlp(img_ffn_input) * img_mod2_gate.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_ffn_input) * txt_mod2_gate.unsqueeze(1) + + return hidden_states, encoder_hidden_states + + +class JoyImageTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding( + in_channels=time_freq_dim, + time_embed_dim=dim, + dtype=dtype, + device=device, + operations=operations, + ) + self.act_fn = nn.SiLU() + self.time_proj = operations.Linear(dim, time_proj_dim, bias=True, dtype=dtype, device=device) + self.text_embedder = PixArtAlphaTextProjection( + text_embed_dim, dim, act_fn="gelu_tanh", dtype=dtype, device=device, operations=operations, + ) + + def forward(self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor): + timestep = self.timesteps_proj(timestep) + temb = self.time_embedder(timestep.to(dtype=encoder_hidden_states.dtype)).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + return temb, timestep_proj, encoder_hidden_states + + +class JoyImageTransformer3DModel(nn.Module): + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 16, + out_channels: Optional[int] = None, + hidden_size: int = 3072, + num_attention_heads: int = 24, + text_dim: int = 4096, + mlp_width_ratio: float = 4.0, + num_layers: int = 20, + rope_dim_list: list = [16, 56, 56], + theta: int = 256, + image_model=None, + dtype=None, + device=None, + operations=None, + ): + super().__init__() + self.dtype = dtype + self.out_channels = out_channels or in_channels + self.patch_size = list(patch_size) + self.rope_dim_list = list(rope_dim_list) + self.theta = theta + + attention_head_dim = hidden_size // num_attention_heads + + self.img_in = operations.Conv3d( + in_channels, + hidden_size, + kernel_size=tuple(self.patch_size), + stride=tuple(self.patch_size), + dtype=dtype, + device=device, + ) + + self.condition_embedder = JoyImageTimeTextImageEmbedding( + dim=hidden_size, + time_freq_dim=256, + time_proj_dim=hidden_size * 6, + text_embed_dim=text_dim, + dtype=dtype, + device=device, + operations=operations, + ) + + self.double_blocks = nn.ModuleList([ + JoyImageTransformerBlock( + dim=hidden_size, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + dtype=dtype, + device=device, + operations=operations, + ) + for _ in range(num_layers) + ]) + + self.norm_out = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) + self.proj_out = operations.Linear( + hidden_size, + self.out_channels * math.prod(self.patch_size), + bias=True, + dtype=dtype, + device=device, + ) + + def _get_rotary_pos_embed_for_range( + self, + start: Tuple[int, int, int], + stop: Tuple[int, int, int], + device=None, + ) -> torch.Tensor: + # 3D RoPE for the patch grid range [start, stop) over (t, h, w). Token order after + # reshape(-1) is (t, h, w), matching the img_in Conv3d flatten. + rope_dim_list = self.rope_dim_list + + grids = [torch.arange(start[i], stop[i], dtype=torch.float32, device=device) for i in range(3)] + mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0) + + angles_parts = [] + for i, dim in enumerate(rope_dim_list): + pos = mesh[i].reshape(-1) + freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device)[: (dim // 2)] / dim)) + angles_parts.append(torch.outer(pos, freqs)) + + angles = torch.cat(angles_parts, dim=1) + cos = angles.cos() + sin = angles.sin() + return torch.stack((cos, -sin, sin, cos), dim=-1).unflatten(-1, (2, 2)) + + def get_rotary_pos_embed_for_components( + self, + component_sizes, + device=None, + ) -> torch.Tensor: + # Per-component 3D RoPE. component_sizes is a list of (t, h, w) patch grid sizes in + # sequence order [target, ref0, ref1, ...]; h/w restart at 0 for each component while t + # continues from the running offset, giving every image its own temporal position band. + freqs_parts = [] + t_offset = 0 + for (t, h, w) in component_sizes: + freqs = self._get_rotary_pos_embed_for_range( + start=(t_offset, 0, 0), + stop=(t_offset + t, h, w), + device=device, + ) + freqs_parts.append(freqs) + t_offset += t + return torch.cat(freqs_parts, dim=0).unsqueeze(0).unsqueeze(2) + + def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + c = self.out_channels + pt, ph, pw = self.patch_size + x = x.reshape(x.shape[0], t, h, w, pt, ph, pw, c) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) + return x.reshape(x.shape[0], c, t * pt, h * ph, w * pw) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor = None, + ref_latents=None, + control=None, + transformer_options=None, + **kwargs, + ) -> torch.Tensor: + transformer_options = {} if transformer_options is None else transformer_options.copy() + return comfy.patcher_extension.WrapperExecutor.new_class_executor( + self._forward, + self, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options) + ).execute(hidden_states, timestep, context, ref_latents, transformer_options, **kwargs) + + def _forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + ref_latents=None, + transformer_options=None, + **kwargs, + ) -> torch.Tensor: + pt, ph, pw = self.patch_size + _, _, ot, oh, ow = hidden_states.shape + + components = [hidden_states, *(ref_latents or [])] + component_sizes = [] + img_tokens = [] + for comp in components: + comp = comfy.ldm.common_dit.pad_to_patch_size(comp, self.patch_size) + _, _, ct, ch, cw = comp.shape + component_sizes.append((ct // pt, ch // ph, cw // pw)) + tokens = self.img_in(comp).flatten(2).transpose(1, 2) # (B, n_i, D) + img_tokens.append(tokens) + + img = torch.cat(img_tokens, dim=1) + + _, vec, txt = self.condition_embedder(timestep, context) + vec = vec.unflatten(1, (6, -1)) + + image_rotary_emb = self.get_rotary_pos_embed_for_components( + component_sizes, + device=hidden_states.device, + ) + + patches_replace = transformer_options.get("patches_replace", {}) + blocks_replace = patches_replace.get("dit", {}) + transformer_options["total_blocks"] = len(self.double_blocks) + transformer_options["block_type"] = "double" + for i, block in enumerate(self.double_blocks): + transformer_options["block_index"] = i + if ("double_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["img"], out["txt"] = block( + hidden_states=args["img"], + encoder_hidden_states=args["txt"], + temb=args["vec"], + image_rotary_emb=args["pe"], + transformer_options=args.get("transformer_options"), + ) + return out + + out = blocks_replace[("double_block", i)]({"img": img, + "txt": txt, + "vec": vec, + "pe": image_rotary_emb, + "transformer_options": transformer_options}, + {"original_block": block_wrap}) + txt = out["txt"] + img = out["img"] + else: + img, txt = block( + hidden_states=img, + encoder_hidden_states=txt, + temb=vec, + image_rotary_emb=image_rotary_emb, + transformer_options=transformer_options, + ) + + tt, th, tw = component_sizes[0] + target_tokens = tt * th * tw + img = img[:, :target_tokens, :] + img = self.proj_out(self.norm_out(img)) + img = self.unpatchify(img, tt, th, tw) + return img[:, :, :ot, :oh, :ow] diff --git a/comfy/model_base.py b/comfy/model_base.py index 786a7c127..98f5ba48b 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -58,6 +58,7 @@ import comfy.ldm.omnigen.omnigen2 import comfy.ldm.seedvr.model import comfy.ldm.boogu.model import comfy.ldm.qwen_image.model +import comfy.ldm.joyimage.model import comfy.ldm.ideogram4.model import comfy.ldm.krea2.model import comfy.ldm.kandinsky5.model @@ -2276,6 +2277,28 @@ class QwenImage(BaseModel): out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16]) return out +class JoyImage(BaseModel): + def __init__(self, model_config, model_type=ModelType.FLOW, device=None): + super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel) + self.memory_usage_factor_conds = ("ref_latents",) + + def extra_conds(self, **kwargs): + out = super().extra_conds(**kwargs) + cross_attn = kwargs.get("cross_attn", None) + if cross_attn is not None: + out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents]) + return out + + def extra_conds_shapes(self, **kwargs): + out = {} + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16]) + return out + class Ideogram4(BaseModel): def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 70c8625e3..a1bf047f8 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -1058,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): dit_config["image_model"] = "SAM31" return dit_config + if ( + '{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys + and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys + and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys + and '{}img_in.weight'.format(key_prefix) in state_dict_keys + and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5 + ): + img_in = state_dict['{}img_in.weight'.format(key_prefix)] + head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0] + return { + "image_model": "joyimage", + "in_channels": img_in.shape[1], + "hidden_size": img_in.shape[0], + "patch_size": list(img_in.shape[2:]), + "num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'), + "num_attention_heads": img_in.shape[0] // head_dim, + "text_dim": 4096, + } + if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys: return None diff --git a/comfy/sd.py b/comfy/sd.py index 4a0742e7a..9d7fa731f 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -76,6 +76,7 @@ import comfy.text_encoders.gemma4 import comfy.text_encoders.cogvideo import comfy.text_encoders.sa3 import comfy.text_encoders.gpt_oss +import comfy.text_encoders.joyimage import comfy.model_patcher import comfy.lora @@ -1377,6 +1378,7 @@ class CLIPType(Enum): IDEOGRAM4 = 30 BOOGU = 31 KREA2 = 32 + JOYIMAGE = 33 @@ -1706,6 +1708,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."}) clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data)) clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer + elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx. + clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."}) + clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data)) + clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused. klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b" clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index b82e4178f..e7c8983aa 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -27,6 +27,7 @@ import comfy.text_encoders.z_image import comfy.text_encoders.ideogram4 import comfy.text_encoders.boogu import comfy.text_encoders.krea2 +import comfy.text_encoders.joyimage import comfy.text_encoders.anima import comfy.text_encoders.ace15 import comfy.text_encoders.longcat_image @@ -1911,6 +1912,38 @@ class QwenImage(supported_models_base.BASE): hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref)) return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect)) +class JoyImage(supported_models_base.BASE): + unet_config = { + "image_model": "joyimage", + } + + sampling_settings = { + "multiplier": 1000, + "shift": 1.5, + } + + memory_usage_factor = 1.8 + + unet_extra_config = { + "theta": 10000, + "rope_dim_list": [16, 56, 56], + } + + latent_format = latent_formats.Wan21 + + supported_inference_dtypes = [torch.bfloat16, torch.float32] + + vae_key_prefix = ["vae."] + text_encoder_key_prefix = ["text_encoders."] + + def get_model(self, state_dict, prefix="", device=None): + return model_base.JoyImage(self, device=device) + + def clip_target(self, state_dict={}): + pref = self.text_encoder_key_prefix[0] + qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl.transformer.".format(pref)) + return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect)) + class HunyuanImage21(HunyuanVideo): unet_config = { "image_model": "hunyuan_video", @@ -2389,6 +2422,7 @@ models = [ Omnigen2, Boogu, QwenImage, + JoyImage, Ideogram4, Krea2, Flux2, diff --git a/comfy/text_encoders/joyimage.py b/comfy/text_encoders/joyimage.py new file mode 100644 index 000000000..143c44250 --- /dev/null +++ b/comfy/text_encoders/joyimage.py @@ -0,0 +1,97 @@ +import torch + +from comfy import sd1_clip +import comfy.text_encoders.qwen_vl +from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer + +JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" +JOYIMAGE_TEMPLATE_TEXT = ( + "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|im_end|>\n" + "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" +) +JOYIMAGE_TEMPLATE_IMAGE = ( + "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|im_end|>\n" + f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n" +) +# The DiT was trained without the leading system-prompt tokens. +JOYIMAGE_DROP_IDX = 34 +PAD_TOKEN = 151643 + + +class Qwen3VL8B_JoyImage(Qwen3VL): + model_type = "qwen3vl_8b" + + def preprocess_embed(self, embed, device): + if embed["type"] == "image": + image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images( + embed["data"], min_pixels=65536, max_pixels=16777216, patch_size=16, + image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5], + interpolation="bicubic", + ) + merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid) + return merged, {"grid": grid, "deepstack": deepstack} + return None, None + + +class JoyImageTokenizer(Qwen3VLTokenizer): + def __init__(self, embedding_directory=None, tokenizer_data={}): + super().__init__( + embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, + model_type="qwen3vl_8b", + ) + self.llama_template = JOYIMAGE_TEMPLATE_TEXT + self.llama_template_images = JOYIMAGE_TEMPLATE_IMAGE + + def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=None, **kwargs): + kwargs.pop("thinking", None) + return super().tokenize_with_weights( + text, return_word_ids=return_word_ids, llama_template=llama_template, + images=images or [], thinking=True, **kwargs, + ) + + +class _JoyImageClipModel(sd1_clip.SDClipModel): + def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None, + attention_mask=True, model_options={}): + super().__init__( + device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, + # JoyImage conditions on the pre-final-norm output of the last decoder layer. + dtype=dtype, special_tokens={"pad": PAD_TOKEN}, layer_norm_hidden_state=False, + model_class=Qwen3VL8B_JoyImage, enable_attention_masks=attention_mask, + return_attention_masks=attention_mask, model_options=model_options, + ) + + +class JoyImageTEModel(sd1_clip.SD1ClipModel): + def __init__(self, device="cpu", dtype=None, model_options={}): + super().__init__( + device=device, dtype=dtype, name="qwen3vl_8b", + clip_model=_JoyImageClipModel, model_options=model_options, + ) + + def encode_token_weights(self, token_weight_pairs): + out, pooled, extra = super().encode_token_weights(token_weight_pairs) + if out.shape[1] <= JOYIMAGE_DROP_IDX: + raise ValueError( + f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter " + f"than drop_idx={JOYIMAGE_DROP_IDX}; the prompt did not include the " + f"template prefix." + ) + out = out[:, JOYIMAGE_DROP_IDX:] + if "attention_mask" in extra: + extra["attention_mask"] = extra["attention_mask"][:, JOYIMAGE_DROP_IDX:] + return out, pooled, extra + + +def te(dtype_llama=None, llama_quantization_metadata=None): + class JoyImageTEModel_(JoyImageTEModel): + def __init__(self, device="cpu", dtype=None, model_options={}): + if llama_quantization_metadata is not None: + model_options = model_options.copy() + model_options["quantization_metadata"] = llama_quantization_metadata + if dtype_llama is not None: + dtype = dtype_llama + super().__init__(device=device, dtype=dtype, model_options=model_options) + return JoyImageTEModel_ diff --git a/comfy/text_encoders/qwen_vl.py b/comfy/text_encoders/qwen_vl.py index 924eb6ad8..f97a88061 100644 --- a/comfy/text_encoders/qwen_vl.py +++ b/comfy/text_encoders/qwen_vl.py @@ -15,6 +15,7 @@ def process_qwen2vl_images( merge_size: int = 2, image_mean: list = None, image_std: list = None, + interpolation: str = "bilinear", ): if image_mean is None: image_mean = [0.48145466, 0.4578275, 0.40821073] @@ -47,10 +48,9 @@ def process_qwen2vl_images( img_resized = F.interpolate( img.unsqueeze(0), size=(h_bar, w_bar), - mode='bilinear', + mode=interpolation, align_corners=False ).squeeze(0) - normalized = img_resized.clone() for c in range(3): normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c] diff --git a/comfy_api_nodes/apis/openai.py b/comfy_api_nodes/apis/openai.py index bee75d639..827281788 100644 --- a/comfy_api_nodes/apis/openai.py +++ b/comfy_api_nodes/apis/openai.py @@ -128,7 +128,7 @@ class OpenAIResponse(ModelResponseProperties, ResponseProperties): parallel_tool_calls: bool | None = Field(True) status: str | None = Field( None, - description="One of `completed`, `failed`, `in_progress`, or `incomplete`.", + description="One of `completed`, `failed`, `in_progress`, `incomplete`, `queued`, or `cancelled`.", ) usage: ResponseUsage | None = Field(None) diff --git a/comfy_api_nodes/nodes_openai.py b/comfy_api_nodes/nodes_openai.py index ad62f2164..de2c94353 100644 --- a/comfy_api_nodes/nodes_openai.py +++ b/comfy_api_nodes/nodes_openai.py @@ -41,6 +41,9 @@ STARTING_POINT_ID_PATTERN = r"" class SupportedOpenAIModel(str, Enum): + gpt_5_6_sol = "gpt-5.6-sol" + gpt_5_6_terra = "gpt-5.6-terra" + gpt_5_6_luna = "gpt-5.6-luna" gpt_5_5_pro = "gpt-5.5-pro" gpt_5_5 = "gpt-5.5" gpt_5 = "gpt-5" @@ -1063,6 +1066,21 @@ class OpenAIChatNode(IO.ComfyNode): "usd": [0.002, 0.008], "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" } } + : $contains($m, "gpt-5.6-terra") ? { + "type": "list_usd", + "usd": [0.0025, 0.015], + "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" } + } + : $contains($m, "gpt-5.6-luna") ? { + "type": "list_usd", + "usd": [0.001, 0.006], + "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" } + } + : $contains($m, "gpt-5.6") ? { + "type": "list_usd", + "usd": [0.005, 0.03], + "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" } + } : $contains($m, "gpt-5.5-pro") ? { "type": "list_usd", "usd": [0.03, 0.18], diff --git a/comfy_extras/nodes_joyimage.py b/comfy_extras/nodes_joyimage.py new file mode 100644 index 000000000..539dc44b2 --- /dev/null +++ b/comfy_extras/nodes_joyimage.py @@ -0,0 +1,102 @@ +from typing_extensions import override + +import comfy.utils +import node_helpers +from comfy_api.latest import ComfyExtension, io + + +# fmt: off +BUCKETS_1024 = [ + (512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048), + (576, 1600), (576, 1664), (576, 1728), (576, 1792), + (640, 1472), (640, 1536), (640, 1600), + (704, 1344), (704, 1408), (704, 1472), + (768, 1216), (768, 1280), (768, 1344), + (832, 1152), (832, 1216), + (896, 1088), (896, 1152), + (960, 1024), (960, 1088), + (1024, 960), (1024, 1024), + (1088, 896), (1088, 960), + (1152, 832), (1152, 896), + (1216, 768), (1216, 832), + (1280, 768), + (1344, 704), (1344, 768), + (1408, 704), + (1472, 640), (1472, 704), + (1536, 640), + (1600, 576), (1600, 640), + (1664, 576), + (1728, 576), + (1792, 512), (1792, 576), + (1856, 512), + (1920, 512), + (1984, 512), + (2048, 512), +] +# fmt: on + + +def _find_best_bucket(height: int, width: int) -> tuple[int, int]: + target_ratio = height / width + return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio)) + + +def _resize_reference(image): + if image.shape[0] != 1: + raise ValueError("JoyImage reference inputs must contain one image each") + samples = image.movedim(-1, 1) + bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3]) + resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") + return resized.movedim(1, -1)[:, :, :, :3] + + +def _encode(clip, prompt, vae, images): + resized_images = [_resize_reference(image) for image in images] + conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images)) + if vae is not None and resized_images: + ref_latents = [vae.encode(image) for image in resized_images] + conditioning = node_helpers.conditioning_set_values( + conditioning, {"reference_latents": ref_latents}, append=True, + ) + return conditioning + + +class TextEncodeJoyImageEdit(io.ComfyNode): + @classmethod + def define_schema(cls): + image_template = io.Autogrow.TemplatePrefix( + io.Image.Input("image"), + prefix="image", + min=0, + max=6, + ) + return io.Schema( + node_id="TextEncodeJoyImageEdit", + category="model/conditioning/joyimage", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("prompt", multiline=True, dynamic_prompts=True), + io.Vae.Input("vae", optional=True), + io.Autogrow.Input("images", template=image_template, optional=True), + ], + outputs=[ + io.Conditioning.Output(), + ], + ) + + @classmethod + def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput: + images = images or {} + return io.NodeOutput(_encode(clip, prompt, vae, list(images.values()))) + + +class JoyImageExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + TextEncodeJoyImageEdit, + ] + + +async def comfy_entrypoint() -> JoyImageExtension: + return JoyImageExtension() diff --git a/nodes.py b/nodes.py index 883258bd1..b03d6c603 100644 --- a/nodes.py +++ b/nodes.py @@ -992,7 +992,7 @@ class CLIPLoader: @classmethod def INPUT_TYPES(s): return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ), - "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ), + "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage"], ), }, "optional": { "device": (["default", "cpu"], {"advanced": True}), @@ -1002,7 +1002,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\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" def load_clip(self, clip_name, type="stable_diffusion", device="default"): clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) @@ -2462,6 +2462,7 @@ async def init_builtin_extra_nodes(): "nodes_seedvr.py", "nodes_context_windows.py", "nodes_qwen.py", + "nodes_joyimage.py", "nodes_boogu.py", "nodes_chroma_radiance.py", "nodes_pid.py", diff --git a/tests-unit/comfy_test/model_detection_test.py b/tests-unit/comfy_test/model_detection_test.py index 7c5b271c5..b40ea0d4c 100644 --- a/tests-unit/comfy_test/model_detection_test.py +++ b/tests-unit/comfy_test/model_detection_test.py @@ -112,6 +112,17 @@ def _make_pid_v1_5_sd(latent_proj_channels=16): return sd +def _make_joyimage_edit_plus_sd(): + sd = { + "img_in.weight": torch.empty(4096, 16, 1, 2, 2, device="meta"), + "condition_embedder.time_embedder.linear_1.weight": torch.empty(1, device="meta"), + "double_blocks.0.attn.img_attn_q_norm.weight": torch.empty(128, device="meta"), + } + for i in range(40): + sd[f"double_blocks.{i}.attn.img_attn_qkv.weight"] = torch.empty(1, device="meta") + return sd + + def _add_model_diffusion_prefix(sd): return {f"model.diffusion_model.{k}": v for k, v in sd.items()} @@ -258,6 +269,26 @@ class TestModelDetection: assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,) assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,) + def test_joyimage_edit_plus_detection(self): + sd = _make_joyimage_edit_plus_sd() + unet_config = detect_unet_config(sd, "") + + assert unet_config == { + "image_model": "joyimage", + "in_channels": 16, + "hidden_size": 4096, + "patch_size": [1, 2, 2], + "num_layers": 40, + "num_attention_heads": 32, + "text_dim": 4096, + } + assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "JoyImage" + + def test_incomplete_joyimage_signature_is_not_detected(self): + sd = _make_joyimage_edit_plus_sd() + del sd["double_blocks.0.attn.img_attn_q_norm.weight"] + assert detect_unet_config(sd, "") is None + def test_unet_config_and_required_keys_combination_is_unique(self): """Each model in the registry must have a unique combination of ``unet_config`` and ``required_keys``. If two models share the same