diff --git a/blueprints/Text to Image (LongCat-Image).json b/blueprints/Text to Image (LongCat-Image).json new file mode 100644 index 000000000..36b021359 --- /dev/null +++ b/blueprints/Text to Image (LongCat-Image).json @@ -0,0 +1 @@ +{"id": "a7e3b1c0-4f2d-4e8a-9b1c-longcat00001", "revision": 0, "last_node_id": 20, "last_link_id": 20, "nodes": [{"id": 1, "type": "lc-subgraph-001", "pos": [0, 1230], "size": [400, 470], "flags": {}, "order": 0, "mode": 0, "inputs": [{"label": "prompt", "name": "text", "type": "STRING", "widget": {"name": "text"}, "link": null}, {"name": "width", "type": "INT", "widget": {"name": "width"}, "link": null}, {"name": "height", "type": "INT", "widget": {"name": "height"}, "link": null}, {"name": "unet_name", "type": "COMBO", "widget": {"name": "unet_name"}, "link": null}, {"name": "clip_name", "type": "COMBO", "widget": {"name": "clip_name"}, "link": null}, {"name": "vae_name", "type": "COMBO", "widget": {"name": "vae_name"}, "link": null}], "outputs": [{"localized_name": "IMAGE", "name": "IMAGE", "type": "IMAGE", "links": []}], "properties": {"proxyWidgets": [["-1", "text"], ["-1", "width"], ["-1", "height"], ["3", "seed"], ["3", "control_after_generate"], ["-1", "unet_name"], ["-1", "clip_name"], ["-1", "vae_name"]], "cnr_id": "comfy-core", "ver": "0.3.73", "enableTabs": false}, "widgets_values": ["A young Asian woman wearing a yellow knit sweater with a white necklace, sitting with her hands on her knees and a serene expression. The background is a rough brick wall with warm afternoon sunlight.", 768, 1344, null, null, "", "", "ae.safetensors"]}], "links": [], "groups": [], "definitions": {"subgraphs": [{"id": "lc-subgraph-001", "version": 1, "state": {"lastGroupId": 4, "lastNodeId": 20, "lastLinkId": 20, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "local-Text to Image (LongCat-Image)", "inputNode": {"id": -10, "bounding": [-80, 425, 120, 160]}, "outputNode": {"id": -20, "bounding": [1490, 415, 120, 60]}, "inputs": [{"id": "inp-text", "name": "text", "type": "STRING", "linkIds": [10], "label": "prompt", "pos": [20, 445]}, {"id": "inp-width", "name": "width", "type": "INT", "linkIds": [11], "pos": [20, 465]}, {"id": "inp-height", "name": "height", "type": "INT", "linkIds": [12], "pos": [20, 485]}, {"id": "inp-unet", "name": "unet_name", "type": "COMBO", "linkIds": [13], "pos": [20, 505]}, {"id": "inp-clip", "name": "clip_name", "type": "COMBO", "linkIds": [14], "pos": [20, 525]}, {"id": "inp-vae", "name": "vae_name", "type": "COMBO", "linkIds": [15], "pos": [20, 545]}], "outputs": [{"id": "out-image", "name": "IMAGE", "type": "IMAGE", "linkIds": [9], "localized_name": "IMAGE", "pos": [1510, 435]}], "widgets": [], "nodes": [{"id": 1, "type": "UNETLoader", "pos": [110, 200], "size": [270, 82], "flags": {}, "order": 0, "mode": 0, "inputs": [{"name": "unet_name", "type": "COMBO", "widget": {"name": "unet_name"}, "link": 13}, {"name": "weight_dtype", "type": "COMBO", "widget": {"name": "weight_dtype"}, "link": null}], "outputs": [{"name": "MODEL", "type": "MODEL", "links": [1]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "UNETLoader"}, "widgets_values": ["", "default"]}, {"id": 2, "type": "CLIPLoader", "pos": [110, 330], "size": [270, 106], "flags": {}, "order": 1, "mode": 0, "inputs": [{"name": "clip_name", "type": "COMBO", "widget": {"name": "clip_name"}, "link": 14}, {"name": "type", "type": "COMBO", "widget": {"name": "type"}, "link": null}, {"name": "device", "shape": 7, "type": "COMBO", "widget": {"name": "device"}, "link": null}], "outputs": [{"name": "CLIP", "type": "CLIP", "links": [2]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "CLIPLoader"}, "widgets_values": ["", "longcat_image", "default"]}, {"id": 3, "type": "VAELoader", "pos": [110, 480], "size": [270, 58], "flags": {}, "order": 2, "mode": 0, "inputs": [{"name": "vae_name", "type": "COMBO", "widget": {"name": "vae_name"}, "link": 15}], "outputs": [{"name": "VAE", "type": "VAE", "links": [3]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "VAELoader", "models": [{"name": "ae.safetensors", "url": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors", "directory": "vae"}]}, "widgets_values": ["ae.safetensors"]}, {"id": 4, "type": "CLIPTextEncodeLongCatImage", "pos": [430, 200], "size": [410, 320], "flags": {}, "order": 3, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 2}, {"name": "text", "type": "STRING", "widget": {"name": "text"}, "link": 10}, {"name": "guidance", "type": "FLOAT", "widget": {"name": "guidance"}, "link": null}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [4, 5]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "CLIPTextEncodeLongCatImage"}, "widgets_values": ["", 4.0]}, {"id": 5, "type": "ConditioningZeroOut", "pos": [640, 620], "size": [204, 26], "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "conditioning", "type": "CONDITIONING", "link": 5}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [6]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "ConditioningZeroOut"}, "widgets_values": []}, {"id": 6, "type": "EmptySD3LatentImage", "pos": [110, 630], "size": [260, 106], "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "width", "type": "INT", "widget": {"name": "width"}, "link": 11}, {"name": "height", "type": "INT", "widget": {"name": "height"}, "link": 12}, {"name": "batch_size", "type": "INT", "widget": {"name": "batch_size"}, "link": null}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [7]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "EmptySD3LatentImage"}, "widgets_values": [768, 1344, 1]}, {"id": 7, "type": "KSampler", "pos": [880, 270], "size": [315, 262], "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 1}, {"name": "positive", "type": "CONDITIONING", "link": 4}, {"name": "negative", "type": "CONDITIONING", "link": 6}, {"name": "latent_image", "type": "LATENT", "link": 7}, {"name": "seed", "type": "INT", "widget": {"name": "seed"}, "link": null}, {"name": "steps", "type": "INT", "widget": {"name": "steps"}, "link": null}, {"name": "cfg", "type": "FLOAT", "widget": {"name": "cfg"}, "link": null}, {"name": "sampler_name", "type": "COMBO", "widget": {"name": "sampler_name"}, "link": null}, {"name": "scheduler", "type": "COMBO", "widget": {"name": "scheduler"}, "link": null}, {"name": "denoise", "type": "FLOAT", "widget": {"name": "denoise"}, "link": null}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [8]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "KSampler"}, "widgets_values": [0, "randomize", 50, 4.0, "euler", "simple", 1]}, {"id": 8, "type": "VAEDecode", "pos": [1220, 160], "size": [210, 46], "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 8}, {"name": "vae", "type": "VAE", "link": 3}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [9]}], "properties": {"cnr_id": "comfy-core", "ver": "0.3.73", "Node name for S&R": "VAEDecode"}, "widgets_values": []}], "groups": [{"id": 1, "title": "Image size", "bounding": [100, 560, 290, 200], "color": "#3f789e", "font_size": 24, "flags": {}}, {"id": 2, "title": "Prompt", "bounding": [410, 130, 450, 540], "color": "#3f789e", "font_size": 24, "flags": {}}, {"id": 3, "title": "Models", "bounding": [100, 130, 290, 413], "color": "#3f789e", "font_size": 24, "flags": {}}], "links": [{"id": 1, "origin_id": 1, "origin_slot": 0, "target_id": 7, "target_slot": 0, "type": "MODEL"}, {"id": 2, "origin_id": 2, "origin_slot": 0, "target_id": 4, "target_slot": 0, "type": "CLIP"}, {"id": 3, "origin_id": 3, "origin_slot": 0, "target_id": 8, "target_slot": 1, "type": "VAE"}, {"id": 4, "origin_id": 4, "origin_slot": 0, "target_id": 7, "target_slot": 1, "type": "CONDITIONING"}, {"id": 5, "origin_id": 4, "origin_slot": 0, "target_id": 5, "target_slot": 0, "type": "CONDITIONING"}, {"id": 6, "origin_id": 5, "origin_slot": 0, "target_id": 7, "target_slot": 2, "type": "CONDITIONING"}, {"id": 7, "origin_id": 6, "origin_slot": 0, "target_id": 7, "target_slot": 3, "type": "LATENT"}, {"id": 8, "origin_id": 7, "origin_slot": 0, "target_id": 8, "target_slot": 0, "type": "LATENT"}, {"id": 9, "origin_id": 8, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "IMAGE"}, {"id": 10, "origin_id": -10, "origin_slot": 0, "target_id": 4, "target_slot": 1, "type": "STRING"}, {"id": 11, "origin_id": -10, "origin_slot": 1, "target_id": 6, "target_slot": 0, "type": "INT"}, {"id": 12, "origin_id": -10, "origin_slot": 2, "target_id": 6, "target_slot": 1, "type": "INT"}, {"id": 13, "origin_id": -10, "origin_slot": 3, "target_id": 1, "target_slot": 0, "type": "COMBO"}, {"id": 14, "origin_id": -10, "origin_slot": 4, "target_id": 2, "target_slot": 0, "type": "COMBO"}, {"id": 15, "origin_id": -10, "origin_slot": 5, "target_id": 3, "target_slot": 0, "type": "COMBO"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Image generation and editing/Text to image"}]}, "config": {}, "extra": {"frontendVersion": "1.37.10", "workflowRendererVersion": "LG"}, "version": 0.4} diff --git a/comfy/model_base.py b/comfy/model_base.py index 8f852e3c6..85cd30bae 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -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) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 030ae6980..58fe514c4 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -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" diff --git a/comfy/sd.py b/comfy/sd.py index de119eb8e..7713d4678 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -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 diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 6d08ff0a5..2431a3fb3 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -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] diff --git a/comfy/text_encoders/longcat_image.py b/comfy/text_encoders/longcat_image.py new file mode 100644 index 000000000..84e4860be --- /dev/null +++ b/comfy/text_encoders/longcat_image.py @@ -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_ diff --git a/comfy_extras/nodes_longcat_image.py b/comfy_extras/nodes_longcat_image.py new file mode 100644 index 000000000..5541c5ee3 --- /dev/null +++ b/comfy_extras/nodes_longcat_image.py @@ -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() diff --git a/nodes.py b/nodes.py index bff073e30..adb2faa79 100644 --- a/nodes.py +++ b/nodes.py @@ -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",