Cleanup the joyai PR.
This commit is contained in:
comfyanonymous
parent
864ed49242
commit
a64984fcbe
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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 ``<think>`` 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 "
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
2
nodes.py
2
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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue