diff --git a/comfy/ldm/joyimage/model.py b/comfy/ldm/joyimage/model.py index f59b43908..4e62c9e9c 100644 --- a/comfy/ldm/joyimage/model.py +++ b/comfy/ldm/joyimage/model.py @@ -4,11 +4,11 @@ from typing import Optional, Tuple import torch import torch.nn as nn -import torch.nn.functional as F +import comfy.ldm.common_dit import comfy.ops import comfy.patcher_extension -from comfy.ldm.lightricks.model import TimestepEmbedding, Timesteps +from comfy.ldm.lightricks.model import GELU_approx, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from comfy.ldm.modules.attention import optimized_attention @@ -57,7 +57,7 @@ class JoyImageFeedForward(nn.Module): ): super().__init__() self.net = nn.ModuleList([ - _GeluApproximate(dim, inner_dim, dtype=dtype, device=device, operations=operations), + GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations), nn.Identity(), operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device), ]) @@ -68,15 +68,6 @@ class JoyImageFeedForward(nn.Module): return x -class _GeluApproximate(nn.Module): - def __init__(self, dim_in: int, dim_out: int, dtype=None, device=None, operations=None): - super().__init__() - self.proj = operations.Linear(dim_in, dim_out, bias=True, dtype=dtype, device=device) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return F.gelu(self.proj(x), approximate="tanh") - - class JoyImageAttention(nn.Module): def __init__( self, @@ -106,8 +97,8 @@ class JoyImageAttention(nn.Module): self, img: torch.Tensor, txt: torch.Tensor, - image_rotary_emb: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]], - transformer_options={}, + image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], + transformer_options=None, ) -> Tuple[torch.Tensor, torch.Tensor]: heads = self.num_attention_heads @@ -126,12 +117,7 @@ class JoyImageAttention(nn.Module): txt_q = self.txt_attn_q_norm(txt_q) txt_k = self.txt_attn_k_norm(txt_k) - if image_rotary_emb is not None: - vis_freqs, txt_freqs = image_rotary_emb - if vis_freqs is not None: - img_q, img_k = _apply_rotary_emb(img_q, img_k, vis_freqs) - if txt_freqs is not None: - txt_q, txt_k = _apply_rotary_emb(txt_q, txt_k, txt_freqs) + img_q, img_k = _apply_rotary_emb(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) @@ -142,7 +128,6 @@ class JoyImageAttention(nn.Module): joint_v = joint_v.flatten(2, 3) joint_out = optimized_attention(joint_q, joint_k, joint_v, heads=heads, transformer_options=transformer_options) - joint_out = joint_out.to(joint_q.dtype) seq_img = img.shape[1] img_out = joint_out[:, :seq_img, :] @@ -166,9 +151,6 @@ class JoyImageTransformerBlock(nn.Module): operations=None, ): super().__init__() - self.dim = dim - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim mlp_hidden_dim = int(dim * mlp_width_ratio) self.img_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device) @@ -196,8 +178,8 @@ class JoyImageTransformerBlock(nn.Module): hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None, - transformer_options={}, + image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], + transformer_options=None, ) -> Tuple[torch.Tensor, torch.Tensor]: ( img_mod1_shift, @@ -258,8 +240,8 @@ class JoyImageTimeTextImageEmbedding(nn.Module): ) 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, dtype=dtype, device=device, operations=operations, + 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): @@ -270,17 +252,6 @@ class JoyImageTimeTextImageEmbedding(nn.Module): return temb, timestep_proj, encoder_hidden_states -class _PixArtAlphaTextProjection(nn.Module): - def __init__(self, in_features: int, hidden_size: int, dtype=None, device=None, operations=None): - super().__init__() - self.linear_1 = operations.Linear(in_features, hidden_size, bias=True, dtype=dtype, device=device) - self.act_1 = nn.GELU(approximate="tanh") - self.linear_2 = operations.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device) - - def forward(self, caption: torch.Tensor) -> torch.Tensor: - return self.linear_2(self.act_1(self.linear_1(caption))) - - class JoyImageTransformer3DModel(nn.Module): def __init__( self, @@ -293,7 +264,6 @@ class JoyImageTransformer3DModel(nn.Module): mlp_width_ratio: float = 4.0, num_layers: int = 20, rope_dim_list: list = [16, 56, 56], - rope_type: str = "rope", theta: int = 256, image_model=None, dtype=None, @@ -304,21 +274,10 @@ class JoyImageTransformer3DModel(nn.Module): self.dtype = dtype self.out_channels = out_channels or in_channels self.patch_size = list(patch_size) - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads self.rope_dim_list = list(rope_dim_list) - self.rope_type = rope_type self.theta = theta - if hidden_size % num_attention_heads != 0: - raise ValueError( - f"hidden_size ({hidden_size}) must be divisible by num_attention_heads ({num_attention_heads})" - ) attention_head_dim = hidden_size // num_attention_heads - if sum(self.rope_dim_list) != attention_head_dim: - raise ValueError( - f"sum(rope_dim_list) ({sum(self.rope_dim_list)}) must equal head_dim ({attention_head_dim})" - ) self.img_in = operations.Conv3d( in_channels, @@ -369,12 +328,7 @@ class JoyImageTransformer3DModel(nn.Module): ) -> Tuple[torch.Tensor, 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. - head_dim = self.hidden_size // self.num_attention_heads rope_dim_list = self.rope_dim_list - if rope_dim_list is None: - rope_dim_list = [head_dim // 3 for _ in range(3)] - if sum(rope_dim_list) != head_dim: - raise ValueError("sum(rope_dim_list) should equal head_dim") 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) @@ -413,8 +367,6 @@ class JoyImageTransformer3DModel(nn.Module): def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: c = self.out_channels pt, ph, pw = self.patch_size - if t * h * w != x.shape[1]: - raise ValueError(f"Expected t*h*w ({t * h * w}) to equal x.shape[1] ({x.shape[1]})") 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) @@ -423,75 +375,50 @@ class JoyImageTransformer3DModel(nn.Module): self, hidden_states: torch.Tensor, timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, + context: torch.Tensor = None, ref_latents=None, - transformer_options={}, + 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, encoder_hidden_states, ref_latents, transformer_options, **kwargs) + ).execute(hidden_states, timestep, context, ref_latents, transformer_options, **kwargs) def _forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, + context: torch.Tensor, ref_latents=None, - transformer_options={}, + transformer_options=None, **kwargs, ) -> torch.Tensor: - # The target noise latent and each reference latent are independently patchified by img_in - # (Conv3d) and concatenated along the sequence dim, in the order [target, ref0, ref1, ...]. - # RoPE is built per component so references may differ in resolution. Only the leading - # target segment (tt*th*tw tokens) is projected back out; reference tokens are dropped. - # A single reference is simply the len(ref_latents) == 1 case. - if hidden_states.ndim != 5: - raise ValueError(f"JoyImage transformer expects 5D (B,C,T,H,W) hidden_states; got shape {tuple(hidden_states.shape)}") - - _, _, ot, oh, ow = hidden_states.shape pt, ph, pw = self.patch_size - if ot % pt != 0 or oh % ph != 0 or ow % pw != 0: - raise ValueError( - f"JoyImage: target latent spatial/temporal shape {(ot, oh, ow)} must be divisible by patch_size {tuple(self.patch_size)}" - ) - tt = ot // pt - th = oh // ph - tw = ow // pw - - components = [hidden_states] - if ref_latents is not None: - for r in ref_latents: - if r.ndim != 5: - raise ValueError(f"JoyImage: each reference latent must be 5D (B,C,T,H,W); got shape {tuple(r.shape)}") - components.append(r) + _, _, 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 - if ct % pt != 0 or ch % ph != 0 or cw % pw != 0: - raise ValueError( - f"JoyImage: component shape {(ct, ch, cw)} must be divisible by patch_size {tuple(self.patch_size)}" - ) 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, encoder_hidden_states) + _, vec, txt = self.condition_embedder(timestep, context) vec = vec.unflatten(1, (6, -1)) - vis_cos, vis_sin = self.get_rotary_pos_embed_for_components( + image_rotary_emb = self.get_rotary_pos_embed_for_components( component_sizes, device=hidden_states.device, ) - vis_freqs = (vis_cos, vis_sin) - txt_freqs = None - image_rotary_emb = (vis_freqs, txt_freqs) patches_replace = transformer_options.get("patches_replace", {}) blocks_replace = patches_replace.get("dit", {}) @@ -528,8 +455,9 @@ class JoyImageTransformer3DModel(nn.Module): transformer_options=transformer_options, ) - img = self.proj_out(self.norm_out(img)) + 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 + return img[:, :, :ot, :oh, :ow] diff --git a/comfy/model_base.py b/comfy/model_base.py index 2503ce473..98f5ba48b 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2278,9 +2278,6 @@ class QwenImage(BaseModel): return out class JoyImage(BaseModel): - # The noise latent and every reference latent are concatenated as a token sequence inside the - # transformer. A single-reference edit is just the len(ref_latents) == 1 case. The required CFG - # guidance rescale is applied by the JoyImageGuidanceRescale node (comfy_extras/nodes_joyimage.py). 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",) @@ -2291,17 +2288,8 @@ class JoyImage(BaseModel): 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 None or len(ref_latents) == 0: - raise ValueError( - "JoyImageEdit is an edit model: every conditioning (positive AND negative) must carry " - "reference_latents. Wire the same reference image(s) and vae into both the positive and " - "negative TextEncodeJoyImageEdit / TextEncodeJoyImageEditPlus nodes. Empty negative " - "prompts still need the image(s) and vae." - ) - latents = [] - for lat in ref_latents: - latents.append(self.process_latent_in(lat)) - out['ref_latents'] = comfy.conds.CONDList(latents) + 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): @@ -2311,61 +2299,6 @@ class JoyImage(BaseModel): out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16]) return out - def _apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs): - # Pass the noise latent and the reference latents to the transformer, which patchifies each - # component and concatenates them along the sequence dim. References may be any resolution. - if c_concat is not None: - raise ValueError("JoyImage does not support c_concat / noise_concat conditioning") - transformer_options = transformer_options.copy() - sigma = t - xc = self.model_sampling.calculate_input(sigma, x) - context = c_crossattn - dtype = self.get_dtype_inference() - xc = xc.to(dtype) - device = xc.device - t_in = self.model_sampling.timestep(t).float() - if context is not None: - context = comfy.model_management.cast_to_device(context, device, dtype) - - extra_conds = {} - for o in kwargs: - extra = kwargs[o] - if hasattr(extra, "dtype"): - extra = convert_tensor(extra, dtype, device) - elif isinstance(extra, list): - ex = [] - for ext in extra: - ex.append(convert_tensor(ext, dtype, device)) - extra = ex - extra_conds[o] = extra - - ref_latents = extra_conds.pop("ref_latents", None) - if ref_latents is None or len(ref_latents) == 0: - raise ValueError("JoyImageEdit forward requires ref_latents; got none.") - - if xc.ndim != 5: - raise ValueError("JoyImageEdit: noise latent must be 5D (B,C,T,H,W); got shape {}.".format(tuple(xc.shape))) - - refs = [] - for r in ref_latents: - if r.ndim != 5: - raise ValueError( - "JoyImageEdit: each reference latent must be 5D (B,C,T,H,W); got shape {}.".format(tuple(r.shape)) - ) - refs.append(r) - - if control is not None: - raise ValueError("JoyImageEdit: control (ControlNet) is not supported by the transformer.") - - # The transformer's forward signature is (hidden_states, timestep, encoder_hidden_states, - # ref_latents, transformer_options); it does not accept control/other extra_conds. - if extra_conds: - raise ValueError("JoyImageEdit: unexpected extra_conds keys {} reached the transformer.".format(list(extra_conds.keys()))) - - noise_pred = self.diffusion_model(xc, t_in, context, ref_latents=refs, transformer_options=transformer_options) - - return self.model_sampling.calculate_denoised(sigma, noise_pred.float(), x) - 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 689c2394a..a1bf047f8 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -896,27 +896,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): dit_config["default_ref_method"] = "negative_index" return dit_config - # JoyImageEdit: dual-stream double_blocks with img_attn_qkv, a condition_embedder - # time_embedder, and a 5D Conv3d img_in (kernel [1,2,2]). - if ( - '{}double_blocks.0.attn.img_attn_qkv.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)] - dit_config = {} - dit_config["image_model"] = "joyimage" - dit_config["in_channels"] = img_in.shape[1] - dit_config["hidden_size"] = img_in.shape[0] - dit_config["patch_size"] = list(img_in.shape[2:]) - dit_config["num_layers"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.') - head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0] - dit_config["num_attention_heads"] = dit_config["hidden_size"] // head_dim - # text_dim from the text-embedder input projection - dit_config["text_dim"] = state_dict['{}condition_embedder.text_embedder.linear_1.weight'.format(key_prefix)].shape[1] - return dit_config - if '{}embed_image_indicator.weight'.format(key_prefix) in state_dict_keys: # Ideogram 4 dit_config = {} dit_config["image_model"] = "ideogram4" @@ -1079,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/supported_models.py b/comfy/supported_models.py index 9ea17da28..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 @@ -1916,9 +1917,6 @@ class JoyImage(supported_models_base.BASE): "image_model": "joyimage", } - # multiplier=1000: the transformer's time embedding is trained on t in [0,1000]. - # ModelSamplingDiscreteFlow.timestep(sigma)=sigma*multiplier yields that range; the - # multiplier cancels in the sigma table, so it only rescales the timestep value. sampling_settings = { "multiplier": 1000, "shift": 1.5, @@ -1931,7 +1929,7 @@ class JoyImage(supported_models_base.BASE): "rope_dim_list": [16, 56, 56], } - latent_format = latent_formats.Wan21 # AutoencoderKLWan: z_dim=16, scale_factor_spatial=8, scale_factor_temporal=4. + latent_format = latent_formats.Wan21 supported_inference_dtypes = [torch.bfloat16, torch.float32] @@ -1939,13 +1937,9 @@ class JoyImage(supported_models_base.BASE): text_encoder_key_prefix = ["text_encoders."] def get_model(self, state_dict, prefix="", device=None): - out = model_base.JoyImage(self, device=device) - return out + return model_base.JoyImage(self, device=device) def clip_target(self, state_dict={}): - # Imported lazily so this module stays importable without the text-encoder deps loaded; - # the import is only resolved when a checkpoint is actually loaded. - import comfy.text_encoders.joyimage 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)) diff --git a/comfy/text_encoders/joyimage.py b/comfy/text_encoders/joyimage.py index 8717e9acd..143c44250 100644 --- a/comfy/text_encoders/joyimage.py +++ b/comfy/text_encoders/joyimage.py @@ -1,107 +1,41 @@ -"""JoyImageEdit text encoder: a stock Qwen3-VL-8B multimodal stack feeding the -JoyImageEdit DiT, built on `comfy.text_encoders.qwen3vl` with the -JoyImage-specific prompt templates, system-prompt strip, image preprocessing, -and conditioning-path multimodal handling. -""" - import torch from comfy import sd1_clip import comfy.text_encoders.qwen_vl from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer -# Prompt templates for the text-only and image-conditioned modes. The image-conditioned template -# wraps the user text with one `<|vision_start|><|image_pad|><|vision_end|>` block per reference -# image (no separator between blocks); `{vision}` is filled with the N concatenated blocks and -# `{prompt}` with the user text. +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" - "<|im_start|>user\n{vision}{prompt}<|im_end|>\n<|im_start|>assistant\n" + f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n" ) - -# A single vision block; N copies are concatenated to condition on N reference images. -JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" - -# Number of leading template tokens (system prompt + the user block's opening -# `<|im_start|>`) stripped from the encoded output by -# JoyImageTEModel.encode_token_weights, so the kept sequence begins at the -# `user` token. +# The DiT was trained without the leading system-prompt tokens. JOYIMAGE_DROP_IDX = 34 - -# Special-token ids (vocab shared with Qwen2.5 / Qwen3, vocab_size 151936). -IMAGE_PAD_TOKEN = 151655 PAD_TOKEN = 151643 class Qwen3VL8B_JoyImage(Qwen3VL): - """JoyImage Qwen3-VL-8B encoder. - - Stock `qwen3vl_8b` config (text dims 4096 / 36L / 32H / 8 kv; interleaved - 3D MRoPE rope_dims=[24,20,20], rope_theta=5e6; vision 1152/4304, depth 27, - patch_size 16, deepstack_visual_indexes=[8,16,24]). - """ - model_type = "qwen3vl_8b" def preprocess_embed(self, embed, device): - # Run the vision tower with JoyImage's bicubic+clamp preprocessing and - # return ``(merged, {"grid", "deepstack"})``. 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", clamp=True, + interpolation="bicubic", ) merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid) return merged, {"grid": grid, "deepstack": deepstack} return None, None - def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, - intermediate_output=None, final_layer_norm_intermediate=True, - dtype=None, embeds_info=()): - # The conditioning path must build the 3D MRoPE position ids for the - # image-token block and inject the deepstack visual features. - # `build_image_inputs` returns the kwargs the decoder expects: - # (position_ids, visual_pos_masks, deepstack). - if embeds is not None: - position_ids, visual_pos_masks, deepstack = self.build_image_inputs(embeds, embeds_info) - else: - position_ids, visual_pos_masks, deepstack = None, None, None - return self.model( - x, - attention_mask=attention_mask, - embeds=embeds, - num_tokens=num_tokens, - intermediate_output=intermediate_output, - final_layer_norm_intermediate=final_layer_norm_intermediate, - dtype=dtype, - position_ids=position_ids, - deepstack_embeds=deepstack, - visual_pos_masks=visual_pos_masks, - ) - class JoyImageTokenizer(Qwen3VLTokenizer): - """JoyImageEdit tokenizer. - - ``tokenize_with_weights(text, images=[...])`` selects the image-conditioned - template when one or more image tensors are passed, emitting one - ``<|vision_start|><|image_pad|><|vision_end|>`` block per image (N blocks - for N reference images), otherwise the text-only template. Each - ``<|image_pad|>`` token in the formatted prompt is replaced with an - embedding marker so `SDClipModel.process_tokens` routes each image through - `Qwen3VL8B_JoyImage.preprocess_embed`; ``drop_idx=34`` leading template - tokens are stripped downstream by `JoyImageTEModel.encode_token_weights`. - No ```` block is appended. - """ - def __init__(self, embedding_directory=None, tokenizer_data={}): super().__init__( embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, @@ -110,58 +44,20 @@ class JoyImageTokenizer(Qwen3VLTokenizer): 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=[], **kwargs): - if text.startswith("<|im_start|>"): - llama_text = text - elif llama_template is not None: - llama_text = llama_template.format(text) - elif len(images) > 0: - # One vision block per reference image. - vision = JOYIMAGE_VISION_BLOCK * len(images) - llama_text = self.llama_template_images.format(vision=vision, prompt=text) - else: - llama_text = self.llama_template.format(text) - - # Tokenize the already-rendered template via the grandparent - # (SD1Tokenizer); calling `super()` would re-apply the Qwen3VL template. - tokens = sd1_clip.SD1Tokenizer.tokenize_with_weights( - self, llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs, + 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, ) - key_name = next(iter(tokens)) - embed_count = 0 - qwen_tokens = tokens[key_name] - for r in qwen_tokens: - for i in range(len(r)): - if r[i][0] == IMAGE_PAD_TOKEN: - if len(images) > embed_count: - r[i] = ({"type": "image", "data": images[embed_count], - "original_type": "image"},) + r[i][1:] - embed_count += 1 - if embed_count != len(images): - raise ValueError( - f"JoyImageTokenizer: prompt had {embed_count} <|image_pad|> placeholders " - f"but {len(images)} image(s) were supplied. Either pre-format the prompt " - f"with `<|vision_start|><|image_pad|><|vision_end|>` per image or pass an " - f"image-free prompt." - ) - return tokens - class _JoyImageClipModel(sd1_clip.SDClipModel): - """Qwen3-VL multimodal encoder wrapper. - - Conditions on the **pre-final-norm** output of the last decoder layer - (``layer="hidden", layer_idx=-1, layer_norm_hidden_state=False``). The - post-norm ``last_hidden_state`` differs by ~10x in scale and produces broken - DiT outputs, so these flags must not be changed. - """ - 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, @@ -177,8 +73,6 @@ class JoyImageTEModel(sd1_clip.SD1ClipModel): def encode_token_weights(self, token_weight_pairs): out, pooled, extra = super().encode_token_weights(token_weight_pairs) - # Strip the JOYIMAGE_DROP_IDX-token system-prompt prefix from both the - # embedding sequence and the attention mask. if out.shape[1] <= JOYIMAGE_DROP_IDX: raise ValueError( f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter " diff --git a/comfy/text_encoders/qwen_vl.py b/comfy/text_encoders/qwen_vl.py index 3dc698374..f97a88061 100644 --- a/comfy/text_encoders/qwen_vl.py +++ b/comfy/text_encoders/qwen_vl.py @@ -16,7 +16,6 @@ def process_qwen2vl_images( image_mean: list = None, image_std: list = None, interpolation: str = "bilinear", - clamp: bool = False, ): if image_mean is None: image_mean = [0.48145466, 0.4578275, 0.40821073] @@ -52,9 +51,6 @@ def process_qwen2vl_images( mode=interpolation, align_corners=False ).squeeze(0) - if clamp: - img_resized = img_resized.clamp(0.0, 1.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_extras/nodes_joyimage.py b/comfy_extras/nodes_joyimage.py index bc5eb05aa..539dc44b2 100644 --- a/comfy_extras/nodes_joyimage.py +++ b/comfy_extras/nodes_joyimage.py @@ -1,7 +1,7 @@ -import node_helpers -import comfy.utils -import torch from typing_extensions import override + +import comfy.utils +import node_helpers from comfy_api.latest import ComfyExtension, io @@ -41,147 +41,53 @@ def _find_best_bucket(height: int, width: int) -> tuple[int, int]: return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio)) -class TextEncodeJoyImageEdit(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="TextEncodeJoyImageEdit", - category="advanced/conditioning", - inputs=[ - io.Clip.Input("clip"), - io.String.Input("prompt", multiline=True, dynamic_prompts=True), - io.Vae.Input("vae"), - io.Image.Input("image"), - ], - outputs=[ - io.Conditioning.Output(), - io.Image.Output(display_name="image"), - ], - ) - - @classmethod - def execute(cls, clip, prompt, vae, image) -> io.NodeOutput: - samples = image.movedim(-1, 1) - src_h, src_w = samples.shape[2], samples.shape[3] - bucket_h, bucket_w = _find_best_bucket(src_h, src_w) - - resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") - resized_image = resized.movedim(1, -1)[:, :, :, :3] - - tokens = clip.tokenize(prompt, images=[resized_image]) - conditioning = clip.encode_from_tokens_scheduled(tokens) - - ref_latent = vae.encode(resized_image) - conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": [ref_latent]}, append=True) - - # Return the bucketed image so VAEEncode can build a matching-size init latent; - # the bucket size isn't reproducible outside this node. - return io.NodeOutput(conditioning, resized_image) +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] -class TextEncodeJoyImageEditPlus(io.ComfyNode): - """JoyImageEdit multi-image (Plus) text-encode node. - - Accepts 1-6 optional reference images. Each supplied image is - bucket-resized independently (same buckets/resize as the single-image - node), VAE-encoded, and appended in order to - ``conditioning["reference_latents"]`` (image1 → ref0, image2 → ref1, ...). - All resized images are passed to the VL tower in one call; the tokenizer - emits one ``<|vision_start|><|image_pad|><|vision_end|>`` block per image. - """ - - MAX_IMAGES = 6 - - @classmethod - def define_schema(cls): - return io.Schema( - node_id="TextEncodeJoyImageEditPlus", - category="advanced/conditioning", - inputs=[ - io.Clip.Input("clip"), - io.String.Input("prompt", multiline=True, dynamic_prompts=True), - io.Vae.Input("vae"), - io.Image.Input("image1", optional=True), - io.Image.Input("image2", optional=True), - io.Image.Input("image3", optional=True), - io.Image.Input("image4", optional=True), - io.Image.Input("image5", optional=True), - io.Image.Input("image6", optional=True), - ], - outputs=[ - io.Conditioning.Output(), - io.Image.Output(display_name="image"), - ], - ) - - @classmethod - def execute(cls, clip, prompt, vae, image1=None, image2=None, image3=None, - image4=None, image5=None, image6=None) -> io.NodeOutput: - images = [image1, image2, image3, image4, image5, image6] - supplied = [img for img in images if img is not None] - if len(supplied) == 0: - raise ValueError( - "TextEncodeJoyImageEditPlus requires at least one reference image." - ) - - resized_images = [] - ref_latents = [] - for image in supplied: - samples = image.movedim(-1, 1) - src_h, src_w = samples.shape[2], samples.shape[3] - bucket_h, bucket_w = _find_best_bucket(src_h, src_w) - - resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") - resized_image = resized.movedim(1, -1)[:, :, :, :3] - resized_images.append(resized_image) - ref_latents.append(vae.encode(resized_image)) - - tokens = clip.tokenize(prompt, images=resized_images) - conditioning = clip.encode_from_tokens_scheduled(tokens) +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, ) - - # The last reference sets the target resolution; return it for VAEEncode - # since the bucket size isn't reproducible outside this node. - return io.NodeOutput(conditioning, resized_images[-1]) + return conditioning -class JoyImageGuidanceRescale(io.ComfyNode): - """CFG combine + per-token L2 norm rescale required by JoyImageEdit. - - Wire this onto the model before sampling: JoyImageEdit's diffusers pipeline - rescales the combined noise prediction back to the conditional branch's norm - (comb * ||cond|| / ||comb||), the same rescale CFGNorm's pre_cfg branch does. - """ - +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="JoyImageGuidanceRescale", - category="model/patch", + node_id="TextEncodeJoyImageEdit", + category="model/conditioning/joyimage", inputs=[ - io.Model.Input("model"), + 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.Model.Output(), + io.Conditioning.Output(), ], ) @classmethod - def execute(cls, model) -> io.NodeOutput: - def guidance_rescale(args): - cond = args["cond"] - uncond = args["uncond"] - cond_scale = args["cond_scale"] - comb = uncond + cond_scale * (cond - uncond) - cond_norm = torch.norm(cond, dim=1, keepdim=True) - comb_norm = torch.norm(comb, dim=1, keepdim=True) - return comb * (cond_norm / comb_norm.clamp_min(1e-6)) - - m = model.clone() - m.set_model_sampler_cfg_function(guidance_rescale) - return io.NodeOutput(m) + 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): @@ -189,8 +95,6 @@ class JoyImageExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ TextEncodeJoyImageEdit, - TextEncodeJoyImageEditPlus, - JoyImageGuidanceRescale, ] diff --git a/nodes.py b/nodes.py index 0ea07b8b2..b03d6c603 100644 --- a/nodes.py +++ b/nodes.py @@ -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) 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