Initial commit for LongCat-Image.

This commit is contained in:
Talmaj Marinc 2026-02-22 16:42:18 +01:00
parent 88d05fe483
commit 2e9d4e967b
8 changed files with 383 additions and 2 deletions

File diff suppressed because one or more lines are too long

View File

@ -925,6 +925,25 @@ class Flux(BaseModel):
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()[2:]), ref_latents))])
return out
class LongCatImage(Flux):
def _apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
transformer_options = transformer_options.copy()
rope_opts = transformer_options.get("rope_options", {})
rope_opts = dict(rope_opts)
rope_opts.setdefault("shift_t", 1.0)
rope_opts.setdefault("shift_y", 512.0)
rope_opts.setdefault("shift_x", 512.0)
transformer_options["rope_options"] = rope_opts
return super()._apply_model(x, t, c_concat, c_crossattn, control, transformer_options, **kwargs)
def encode_adm(self, **kwargs):
return None
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
out.pop('guidance', None)
return out
class Flux2(Flux):
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)

View File

@ -282,6 +282,36 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
return dit_config
if '{}x_embedder.weight'.format(key_prefix) in state_dict_keys and '{}transformer_blocks.0.attn.to_q.weight'.format(key_prefix) in state_dict_keys and '{}single_transformer_blocks.0.attn.to_q.weight'.format(key_prefix) in state_dict_keys: #LongCat-Image (diffusers format, Flux variant)
dit_config = {}
dit_config["image_model"] = "flux"
dit_config["axes_dim"] = [16, 56, 56]
dit_config["theta"] = 10000
dit_config["qkv_bias"] = True
dit_config["txt_ids_dims"] = [1, 2]
w = state_dict['{}x_embedder.weight'.format(key_prefix)]
dit_config["hidden_size"] = w.shape[0]
dit_config["in_channels"] = w.shape[1] // 4
dit_config["out_channels"] = dit_config["in_channels"]
dit_config["patch_size"] = 2
ctx_key = '{}context_embedder.weight'.format(key_prefix)
if ctx_key in state_dict_keys:
dit_config["context_in_dim"] = state_dict[ctx_key].shape[1]
else:
dit_config["context_in_dim"] = 3584
dit_config["vec_in_dim"] = None
dit_config["guidance_embed"] = False
dit_config["mlp_ratio"] = 4.0
dit_config["num_heads"] = dit_config["hidden_size"] // sum(dit_config["axes_dim"])
dit_config["depth"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_transformer_blocks.'.format(key_prefix) + '{}.')
return dit_config
if '{}t5_yproj.weight'.format(key_prefix) in state_dict_keys: #Genmo mochi preview
dit_config = {}
dit_config["image_model"] = "mochi_preview"

View File

@ -60,6 +60,7 @@ import comfy.text_encoders.jina_clip_2
import comfy.text_encoders.newbie
import comfy.text_encoders.anima
import comfy.text_encoders.ace15
import comfy.text_encoders.longcat_image
import comfy.model_patcher
import comfy.lora
@ -1160,6 +1161,7 @@ class CLIPType(Enum):
KANDINSKY5_IMAGE = 23
NEWBIE = 24
FLUX2 = 25
LONGCAT_IMAGE = 26
def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
@ -1372,6 +1374,9 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
if clip_type == CLIPType.HUNYUAN_IMAGE:
clip_target.clip = comfy.text_encoders.hunyuan_image.te(byt5=False, **llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
elif clip_type == CLIPType.LONGCAT_IMAGE:
clip_target.clip = comfy.text_encoders.longcat_image.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.longcat_image.LongCatImageTokenizer
else:
clip_target.clip = comfy.text_encoders.qwen_image.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.qwen_image.QwenImageTokenizer

View File

@ -25,6 +25,7 @@ import comfy.text_encoders.kandinsky5
import comfy.text_encoders.z_image
import comfy.text_encoders.anima
import comfy.text_encoders.ace15
import comfy.text_encoders.longcat_image
from . import supported_models_base
from . import latent_formats
@ -1677,6 +1678,142 @@ class ACEStep15(supported_models_base.BASE):
return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, LTXAV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, WAN21_FlowRVS, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, ACEStep15, Omnigen2, QwenImage, Flux2, Kandinsky5Image, Kandinsky5, Anima]
class LongCatImage(supported_models_base.BASE):
unet_config = {
"image_model": "flux",
"guidance_embed": False,
"vec_in_dim": None,
"context_in_dim": 3584,
"txt_ids_dims": [1, 2],
}
sampling_settings = {
}
unet_extra_config = {}
latent_format = latent_formats.Flux
memory_usage_factor = 2.5
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
def process_unet_state_dict(self, state_dict):
out_sd = {}
double_q, double_k, double_v = {}, {}, {}
double_tq, double_tk, double_tv = {}, {}, {}
single_q, single_k, single_v, single_mlp = {}, {}, {}, {}
for k, v in state_dict.items():
if k.startswith("transformer_blocks."):
idx = k.split(".")[1]
rest = ".".join(k.split(".")[2:])
prefix = "double_blocks.{}.".format(idx)
if rest.startswith("norm1.linear."):
out_sd[prefix + "img_mod.lin." + rest.split(".")[-1]] = v
elif rest.startswith("norm1_context.linear."):
out_sd[prefix + "txt_mod.lin." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_q."):
double_q[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_k."):
double_k[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_v."):
double_v[idx + "." + rest.split(".")[-1]] = v
elif rest == "attn.norm_q.weight":
out_sd[prefix + "img_attn.norm.query_norm.weight"] = v
elif rest == "attn.norm_k.weight":
out_sd[prefix + "img_attn.norm.key_norm.weight"] = v
elif rest.startswith("attn.to_out.0."):
out_sd[prefix + "img_attn.proj." + rest.split(".")[-1]] = v
elif rest.startswith("attn.add_q_proj."):
double_tq[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.add_k_proj."):
double_tk[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.add_v_proj."):
double_tv[idx + "." + rest.split(".")[-1]] = v
elif rest == "attn.norm_added_q.weight":
out_sd[prefix + "txt_attn.norm.query_norm.weight"] = v
elif rest == "attn.norm_added_k.weight":
out_sd[prefix + "txt_attn.norm.key_norm.weight"] = v
elif rest.startswith("attn.to_add_out."):
out_sd[prefix + "txt_attn.proj." + rest.split(".")[-1]] = v
elif rest.startswith("ff.net.0.proj."):
out_sd[prefix + "img_mlp.0." + rest.split(".")[-1]] = v
elif rest.startswith("ff.net.2."):
out_sd[prefix + "img_mlp.2." + rest.split(".")[-1]] = v
elif rest.startswith("ff_context.net.0.proj."):
out_sd[prefix + "txt_mlp.0." + rest.split(".")[-1]] = v
elif rest.startswith("ff_context.net.2."):
out_sd[prefix + "txt_mlp.2." + rest.split(".")[-1]] = v
else:
out_sd["double_blocks.{}.{}".format(idx, rest)] = v
elif k.startswith("single_transformer_blocks."):
idx = k.split(".")[1]
rest = ".".join(k.split(".")[2:])
prefix = "single_blocks.{}.".format(idx)
if rest.startswith("norm.linear."):
out_sd[prefix + "modulation.lin." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_q."):
single_q[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_k."):
single_k[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("attn.to_v."):
single_v[idx + "." + rest.split(".")[-1]] = v
elif rest == "attn.norm_q.weight":
out_sd[prefix + "norm.query_norm.weight"] = v
elif rest == "attn.norm_k.weight":
out_sd[prefix + "norm.key_norm.weight"] = v
elif rest.startswith("proj_mlp."):
single_mlp[idx + "." + rest.split(".")[-1]] = v
elif rest.startswith("proj_out."):
out_sd[prefix + "linear2." + rest.split(".")[-1]] = v
else:
out_sd["single_blocks.{}.{}".format(idx, rest)] = v
elif k == "x_embedder.weight" or k == "x_embedder.bias":
out_sd["img_in." + k.split(".")[-1]] = v
elif k == "context_embedder.weight" or k == "context_embedder.bias":
out_sd["txt_in." + k.split(".")[-1]] = v
elif k.startswith("time_embed.timestep_embedder.linear_1."):
out_sd["time_in.in_layer." + k.split(".")[-1]] = v
elif k.startswith("time_embed.timestep_embedder.linear_2."):
out_sd["time_in.out_layer." + k.split(".")[-1]] = v
elif k.startswith("norm_out.linear."):
out_sd["final_layer.adaLN_modulation.1." + k.split(".")[-1]] = v
elif k == "proj_out.weight" or k == "proj_out.bias":
out_sd["final_layer.linear." + k.split(".")[-1]] = v
else:
out_sd[k] = v
for suffix in ["weight", "bias"]:
for idx in sorted(set(x.split(".")[0] for x in double_q)):
qk = idx + "." + suffix
if qk in double_q and qk in double_k and qk in double_v:
out_sd["double_blocks.{}.img_attn.qkv.{}".format(idx, suffix)] = torch.cat([double_q[qk], double_k[qk], double_v[qk]], dim=0)
if qk in double_tq and qk in double_tk and qk in double_tv:
out_sd["double_blocks.{}.txt_attn.qkv.{}".format(idx, suffix)] = torch.cat([double_tq[qk], double_tk[qk], double_tv[qk]], dim=0)
for idx in sorted(set(x.split(".")[0] for x in single_q)):
qk = idx + "." + suffix
if qk in single_q and qk in single_k and qk in single_v and qk in single_mlp:
out_sd["single_blocks.{}.linear1.{}".format(idx, suffix)] = torch.cat([single_q[qk], single_k[qk], single_v[qk], single_mlp[qk]], dim=0)
return out_sd
def get_model(self, state_dict, prefix="", device=None):
out = model_base.LongCatImage(self, device=device)
return out
def clip_target(self, state_dict={}):
pref = self.text_encoder_key_prefix[0]
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
return supported_models_base.ClipTarget(comfy.text_encoders.longcat_image.LongCatImageTokenizer, comfy.text_encoders.longcat_image.te(**hunyuan_detect))
models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, LTXAV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, WAN21_FlowRVS, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, ACEStep15, Omnigen2, QwenImage, LongCatImage, Flux2, Kandinsky5Image, Kandinsky5, Anima]
models += [SVD_img2vid]

View File

@ -0,0 +1,148 @@
import re
import numbers
import torch
from comfy import sd1_clip
from comfy.text_encoders.qwen_image import Qwen25_7BVLITokenizer, Qwen25_7BVLIModel
QUOTE_PAIRS = [("'", "'"), ('"', '"'), ("\u2018", "\u2019"), ("\u201c", "\u201d")]
QUOTE_PATTERN = "|".join(
[re.escape(q1) + r"[^" + re.escape(q1 + q2) + r"]*?" + re.escape(q2) for q1, q2 in QUOTE_PAIRS]
)
WORD_INTERNAL_QUOTE_RE = re.compile(r"[a-zA-Z]+'[a-zA-Z]+")
def split_quotation(prompt):
matches = WORD_INTERNAL_QUOTE_RE.findall(prompt)
mapping = []
for i, word_src in enumerate(set(matches)):
word_tgt = "longcat_$##$_longcat" * (i + 1)
prompt = prompt.replace(word_src, word_tgt)
mapping.append((word_src, word_tgt))
parts = re.split(f"({QUOTE_PATTERN})", prompt)
result = []
for part in parts:
for word_src, word_tgt in mapping:
part = part.replace(word_tgt, word_src)
if not part:
continue
is_quoted = bool(re.match(QUOTE_PATTERN, part))
result.append((part, is_quoted))
return result
class LongCatImageBaseTokenizer(Qwen25_7BVLITokenizer):
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
parts = split_quotation(text)
all_tokens = []
for part_text, is_quoted in parts:
if is_quoted:
for char in part_text:
ids = self.tokenizer(char, add_special_tokens=False)["input_ids"]
all_tokens.extend(ids)
else:
ids = self.tokenizer(part_text, add_special_tokens=False)["input_ids"]
all_tokens.extend(ids)
max_len = self.max_length if self.max_length < 99999999 else 512
if len(all_tokens) > max_len:
all_tokens = all_tokens[:max_len]
output = [(t, 1.0) for t in all_tokens]
return [output]
class LongCatImageTokenizer(sd1_clip.SD1Tokenizer):
def __init__(self, embedding_directory=None, tokenizer_data={}):
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="qwen25_7b", tokenizer=LongCatImageBaseTokenizer)
self.longcat_template_prefix = "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n"
self.longcat_template_suffix = "<|im_end|>\n<|im_start|>assistant\n"
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
skip_template = False
if text.startswith('<|im_start|>'):
skip_template = True
if text.startswith('<|start_header_id|>'):
skip_template = True
if text == '':
text = ' '
base_tok = getattr(self, "qwen25_7b")
if skip_template:
tokens = super().tokenize_with_weights(text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
else:
prefix_ids = base_tok.tokenizer(self.longcat_template_prefix, add_special_tokens=False)["input_ids"]
suffix_ids = base_tok.tokenizer(self.longcat_template_suffix, add_special_tokens=False)["input_ids"]
prompt_tokens = base_tok.tokenize_with_weights(text, return_word_ids=return_word_ids, **kwargs)
prompt_pairs = prompt_tokens[0]
prefix_pairs = [(t, 1.0) for t in prefix_ids]
suffix_pairs = [(t, 1.0) for t in suffix_ids]
combined = prefix_pairs + prompt_pairs + suffix_pairs
tokens = {"qwen25_7b": [combined]}
return tokens
class LongCatImageTEModel(sd1_clip.SD1ClipModel):
def __init__(self, device="cpu", dtype=None, model_options={}):
super().__init__(device=device, dtype=dtype, name="qwen25_7b", clip_model=Qwen25_7BVLIModel, model_options=model_options)
def encode_token_weights(self, token_weight_pairs, template_end=-1):
out, pooled, extra = super().encode_token_weights(token_weight_pairs)
tok_pairs = token_weight_pairs["qwen25_7b"][0]
count_im_start = 0
if template_end == -1:
for i, v in enumerate(tok_pairs):
elem = v[0]
if not torch.is_tensor(elem):
if isinstance(elem, numbers.Integral):
if elem == 151644 and count_im_start < 2:
template_end = i
count_im_start += 1
if out.shape[1] > (template_end + 3):
if tok_pairs[template_end + 1][0] == 872:
if tok_pairs[template_end + 2][0] == 198:
template_end += 3
suffix_start = None
for i in range(len(tok_pairs) - 1, -1, -1):
elem = tok_pairs[i][0]
if not torch.is_tensor(elem) and isinstance(elem, numbers.Integral):
if elem == 151644:
suffix_start = i
break
out = out[:, template_end:]
if "attention_mask" in extra:
extra["attention_mask"] = extra["attention_mask"][:, template_end:]
if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]):
extra.pop("attention_mask")
if suffix_start is not None:
suffix_len = len(tok_pairs) - suffix_start
if suffix_len > 0 and out.shape[1] > suffix_len:
out = out[:, :-suffix_len]
if "attention_mask" in extra:
extra["attention_mask"] = extra["attention_mask"][:, :-suffix_len]
if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]):
extra.pop("attention_mask")
return out, pooled, extra
def te(dtype_llama=None, llama_quantization_metadata=None):
class LongCatImageTEModel_(LongCatImageTEModel):
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 LongCatImageTEModel_

View File

@ -0,0 +1,40 @@
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
class CLIPTextEncodeLongCatImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CLIPTextEncodeLongCatImage",
display_name="CLIP Text Encode (LongCat-Image)",
category="advanced/conditioning/longcat",
description="Text encoding for LongCat-Image with character-level quoted text support. Wrap text in quotes for accurate text rendering.",
inputs=[
io.Clip.Input("clip"),
io.String.Input("text", multiline=True, dynamic_prompts=True),
io.Float.Input("guidance", default=4.0, min=0.0, max=100.0, step=0.1),
],
outputs=[
io.Conditioning.Output(),
],
)
@classmethod
def execute(cls, clip, text, guidance) -> io.NodeOutput:
tokens = clip.tokenize(text)
return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"guidance": guidance}))
encode = execute
class LongCatImageExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
CLIPTextEncodeLongCatImage,
]
async def comfy_entrypoint() -> LongCatImageExtension:
return LongCatImageExtension()

View File

@ -976,7 +976,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"], ),
"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"], ),
},
"optional": {
"device": (["default", "cpu"], {"advanced": True}),
@ -2429,6 +2429,7 @@ async def init_builtin_extra_nodes():
"nodes_tcfg.py",
"nodes_context_windows.py",
"nodes_qwen.py",
"nodes_longcat_image.py",
"nodes_chroma_radiance.py",
"nodes_model_patch.py",
"nodes_easycache.py",