Native LongCat-Image implementation - #12597
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds LongCat-Image support across the codebase: a new Text-to-Image blueprint JSON; new LongCatImage model class and Flux adjustments in comfy/model_base.py; UNet detection update in comfy/model_detection.py; CLIPType enum and LONGCAT_IMAGE text-encoder loading in comfy/sd.py; new supported model entry in comfy/supported_models.py; a LongCatImage tokenizer/TE implementation in comfy/text_encoders/longcat_image.py; two Comfy nodes and an extension in comfy_extras/nodes_longcat_image.py; CLIPLoader option and extras registration; and unit tests for model detection and conversion. 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests-unit/comfy_test/model_detection_test.py (1)
73-73: Unused variableoriginal_models.
original_modelsis assigned but never referenced. Likely leftover from a manual save/restore approach that was replaced bypatch.object.🧹 Remove unused variable
sd = _make_longcat_diffusers_sd() unet_config = detect_unet_config(sd, "") - original_models = comfy.supported_models.models longcat_cls = comfy.supported_models.LongCatImage🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests-unit/comfy_test/model_detection_test.py` at line 73, Remove the unused local variable original_models assigned from comfy.supported_models.models in the test; since patch.object is handling temporary replacement/restore, delete the assignment to original_models to eliminate the dead code and keep the test clean (look for the assignment to original_models and the reference to comfy.supported_models.models in model_detection_test.py).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy_extras/nodes_longcat_image.py`:
- Around line 61-82: The code assumes H and W are divisible by ps=2 before
reshaping (see noise.reshape(...), cond_packed reshape and renormed), which will
raise on odd spatial sizes; add a defensive check right after B, C, H, W =
denoised.shape to verify H % ps == 0 and W % ps == 0 and raise a clear
ValueError including the offending H/W and ps (or alternatively apply symmetric
padding to x/denoised/cond_denoised to make them divisible by ps before the
pack/unpack operations), then proceed with the existing noise/cond packing,
scaling and renorming logic unchanged.
In `@comfy/text_encoders/longcat_image.py`:
- Around line 70-95: In tokenize_with_weights: avoid letting
base_tok.tokenize_with_weights produce 512-length padding before you add the
LongCat template; call base_tok.tokenize_with_weights with padding disabled
(e.g., disable_padding=True or equivalent kwarg) so prompt_pairs is produced
without pre-padding, then build prefix_pairs, prompt_pairs and suffix_pairs into
combined, and only after combining perform truncation/padding to model length
(use your tokenizer's pad/truncate utility or call the shared super method once
on the final combined token list) so prefix_ids/suffix_ids are not separated by
mid-prompt padding; refer to tokenize_with_weights,
base_tok.tokenize_with_weights, prefix_ids, suffix_ids, prompt_pairs and
combined to locate where to change.
- Around line 102-144: The slice logic in encode_token_weights can use
template_end == -1 (no <|im_start|> found) which makes out = out[:, -1:] (last
token); change encode_token_weights to treat a missing marker by setting
template_end = 0 before slicing (or otherwise avoid negative slice) and only
apply the "+3 newline adjustment" when a real marker was detected (i.e., only
run the tok_pairs[template_end + 1]/[+2] checks if template_end was set from the
loop). Update references in encode_token_weights (template_end, tok_pairs, out,
extra, suffix_start) so the slicing/out = out[:, template_end:] and subsequent
attention_mask adjustments are guarded by the marker presence to avoid
accidentally keeping only the last token.
---
Nitpick comments:
In `@tests-unit/comfy_test/model_detection_test.py`:
- Line 73: Remove the unused local variable original_models assigned from
comfy.supported_models.models in the test; since patch.object is handling
temporary replacement/restore, delete the assignment to original_models to
eliminate the dead code and keep the test clean (look for the assignment to
original_models and the reference to comfy.supported_models.models in
model_detection_test.py).
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📥 Commits
Reviewing files that changed from the base of the PR and between dba2766 and bfd302f788942b04c6bf6d65d07645fb2206ddb8.
📒 Files selected for processing (9)
blueprints/Text to Image (LongCat-Image).jsoncomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/longcat_image.pycomfy_extras/nodes_longcat_image.pynodes.pytests-unit/comfy_test/model_detection_test.py
bfd302f to
81a3792
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests-unit/comfy_test/model_detection_test.py (2)
84-89: Considerisinstanceovertype(...).__name__for class identity assertions.String-based class-name checks will silently pass if the class is renamed or imported under an alias.
♻️ Proposed fix
- assert type(result).__name__ == "LongCatImage", ( - f"Expected LongCatImage with order {label}, got {type(result).__name__}" - ) + assert isinstance(result, comfy.supported_models.LongCatImage), ( + f"Expected LongCatImage with order {label}, got {type(result)}" + )And at line 101:
- assert type(model_config).__name__ == "LongCatImage" + assert isinstance(model_config, comfy.supported_models.LongCatImage)Also applies to: 99-101
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests-unit/comfy_test/model_detection_test.py` around lines 84 - 89, Replace fragile string-based class checks with real type checks: instead of asserting type(result).__name__ == "LongCatImage" use an isinstance assertion against the actual class (e.g., assert isinstance(result, LongCatImage)). Update both occurrences (the assertion around model_config_from_unet_config and the similar check at lines ~99-101) and ensure LongCatImage is imported or referenced from the correct module so the isinstance call resolves.
103-113: Test only verifies key presence/absence, skipping the two non-trivial transforms.The PR description calls out Q/K/V fusion and a scale/shift half-swap in
process_unet_state_dictas the critical parts of the conversion. Neither is exercised here — verifying that, say,transformer_blocks.0.attn.to_q.weight+to_k.weight+to_v.weightare fused into a singledouble_blocks.0.img_attn.qkv.weightwith the right shape, and that thenorm_outscale/shift halves are swapped, would meaningfully increase confidence in the conversion correctness.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests-unit/comfy_test/model_detection_test.py` around lines 103 - 113, The test test_longcat_process_unet_state_dict_converts_keys only checks presence/absence of keys but does not validate the two non-trivial transforms in process_unet_state_dict: Q/K/V fusion and the norm_out scale/shift half-swap. Update the test to build source weights for transformer_blocks.0.attn.to_q/to_k/to_v and the norm_out affine, run converted = model_config.process_unet_state_dict(...), then assert the fused tensor exists at double_blocks.0.img_attn.qkv.weight with the expected concatenated shape and contents (verify slices match original to_q/k/v), and assert norm_out parameters have their scale/shift halves swapped compared to the input; reference test_longcat_process_unet_state_dict_converts_keys, process_unet_state_dict, transformer_blocks.*, attn.to_q/to_k/to_v, double_blocks.0.img_attn.qkv.weight, and norm_out in your assertions.comfy_extras/nodes_longcat_image.py (1)
82-83: Replace deprecatedtorch.normwithtorch.linalg.vector_norm.
torch.normis deprecated and may be removed in a future PyTorch release; its documentation and behavior may be incorrect, and it is no longer actively maintained. The recommended replacement for vector norms istorch.linalg.vector_norm().♻️ Proposed fix
- noise_norm = torch.norm(noise_packed, dim=-1, keepdim=True) - cond_norm = torch.norm(cond_packed, dim=-1, keepdim=True) + noise_norm = torch.linalg.vector_norm(noise_packed, dim=-1, keepdim=True) + cond_norm = torch.linalg.vector_norm(cond_packed, dim=-1, keepdim=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy_extras/nodes_longcat_image.py` around lines 82 - 83, Replace deprecated torch.norm calls computing per-vector norms for noise_packed and cond_packed with torch.linalg.vector_norm; specifically update the expressions that assign noise_norm and cond_norm (currently using torch.norm(noise_packed, dim=-1, keepdim=True) and torch.norm(cond_packed, dim=-1, keepdim=True)) to use torch.linalg.vector_norm(noise_packed, dim=-1, keepdim=True) and torch.linalg.vector_norm(cond_packed, dim=-1, keepdim=True) respectively so behavior and signature remain the same but use the supported API.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/model_detection.py`:
- Around line 285-313: The current three-key detection in model_detection.py
(keys '{}x_embedder.weight', '{}transformer_blocks.0.attn.to_q.weight',
'{}single_transformer_blocks.0.attn.to_q.weight' using key_prefix) misidentifies
vanilla Flux diffusers as LongCat-Image; tighten the condition by requiring a
LongCat-specific key (use the existing ctx_key =
'{}context_embedder.weight'.format(key_prefix)) to be present as an additional
positive constraint before building dit_config, so only state dicts that include
context_embedder.weight are considered LongCat-Image (leave count_blocks usage
and subsequent field population unchanged).
---
Nitpick comments:
In `@comfy_extras/nodes_longcat_image.py`:
- Around line 82-83: Replace deprecated torch.norm calls computing per-vector
norms for noise_packed and cond_packed with torch.linalg.vector_norm;
specifically update the expressions that assign noise_norm and cond_norm
(currently using torch.norm(noise_packed, dim=-1, keepdim=True) and
torch.norm(cond_packed, dim=-1, keepdim=True)) to use
torch.linalg.vector_norm(noise_packed, dim=-1, keepdim=True) and
torch.linalg.vector_norm(cond_packed, dim=-1, keepdim=True) respectively so
behavior and signature remain the same but use the supported API.
In `@tests-unit/comfy_test/model_detection_test.py`:
- Around line 84-89: Replace fragile string-based class checks with real type
checks: instead of asserting type(result).__name__ == "LongCatImage" use an
isinstance assertion against the actual class (e.g., assert isinstance(result,
LongCatImage)). Update both occurrences (the assertion around
model_config_from_unet_config and the similar check at lines ~99-101) and ensure
LongCatImage is imported or referenced from the correct module so the isinstance
call resolves.
- Around line 103-113: The test
test_longcat_process_unet_state_dict_converts_keys only checks presence/absence
of keys but does not validate the two non-trivial transforms in
process_unet_state_dict: Q/K/V fusion and the norm_out scale/shift half-swap.
Update the test to build source weights for
transformer_blocks.0.attn.to_q/to_k/to_v and the norm_out affine, run converted
= model_config.process_unet_state_dict(...), then assert the fused tensor exists
at double_blocks.0.img_attn.qkv.weight with the expected concatenated shape and
contents (verify slices match original to_q/k/v), and assert norm_out parameters
have their scale/shift halves swapped compared to the input; reference
test_longcat_process_unet_state_dict_converts_keys, process_unet_state_dict,
transformer_blocks.*, attn.to_q/to_k/to_v, double_blocks.0.img_attn.qkv.weight,
and norm_out in your assertions.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📥 Commits
Reviewing files that changed from the base of the PR and between bfd302f788942b04c6bf6d65d07645fb2206ddb8 and 81a37925e6513e1685a35f2c926916084980315c.
📒 Files selected for processing (9)
blueprints/Text to Image (LongCat-Image).jsoncomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/longcat_image.pycomfy_extras/nodes_longcat_image.pynodes.pytests-unit/comfy_test/model_detection_test.py
✅ Files skipped from review due to trivial changes (1)
- blueprints/Text to Image (LongCat-Image).json
🚧 Files skipped from review as they are similar to previous changes (1)
- nodes.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/model_detection.py`:
- Around line 827-838: The selection currently adds
len(model_config.required_keys) into specificity even when no state_dict was
provided, biasing picks; modify model_config_from_unet_config so that
required_keys are only counted when a state_dict is present/used (i.e., only add
len(model_config.required_keys) to specificity if state_dict is not None and
thus those keys could be validated), preserving the previous first-match
ordering when state_dict is None; refer to model_config_from_unet_config,
best_specificity, model_config.required_keys and the matches(...) call to
implement this conditional weighting.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📥 Commits
Reviewing files that changed from the base of the PR and between 81a37925e6513e1685a35f2c926916084980315c and 449282d062cfb2859d889c60c9d987b02fdfc2b0.
📒 Files selected for processing (2)
comfy/model_detection.pycomfy/supported_models.py
| prefix = "double_blocks.{}.".format(idx) | ||
|
|
||
| if rest.startswith("norm1.linear."): | ||
| out_sd[prefix + "img_mod.lin." + rest.split(".")[-1]] = v |
There was a problem hiding this comment.
Something we can do is repackage/convert the model weights to a format that is easier to load in comfyui instead of having all this logic here it can be a conversion.py file in the hf repo where we host our converted weights.
There was a problem hiding this comment.
Is there some standard for repackaging that I can follow?
There was a problem hiding this comment.
The standard is to make something that doesn't need to be converted by comfyui internally to work.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
blueprints/Text to Image (LongCat-Image).json (1)
1-1: Expose a negative prompt input in the subgraph.The second
CLIPTextEncodeLongCatImagenode isn’t wired to any input, so users can’t set a negative prompt from the blueprint. Consider adding a subgraph input and link for negative text (or remove the node if it’s intentionally fixed).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@blueprints/Text` to Image (LongCat-Image).json at line 1, The subgraph lacks a user-facing negative prompt because the second CLIPTextEncodeLongCatImage node (node id 5, type CLIPTextEncodeLongCatImage) has its "text" input unlinked; add a subgraph input (e.g., inp-neg-text) for the negative prompt and create a link from that subgraph input to node 5's "text" input so users can supply a negative prompt, or if the node is intentionally unused, remove node id 5 and its links (including the link to KSampler) to avoid confusion.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests-unit/comfy_test/model_detection_test.py`:
- Around line 8-36: The fixture _make_longcat_comfyui_sd allocates very large
tensors (H=3072) causing multi-GB memory usage; change H to a small value (e.g.,
32) while keeping C_CTX at 3584 so txt_in.weight.shape[1] still reflects the
real context dim, and keep other shape formulas (C_IN, C_CTX, and all uses like
"img_in.weight", "txt_in.weight", "time_in.*", "final_layer.*", and blocks in
"double_blocks.*" and "single_blocks.*") unchanged so the detection logic that
reads tensor shapes and key presence continues to work but without large memory
allocations.
---
Nitpick comments:
In `@blueprints/Text` to Image (LongCat-Image).json:
- Line 1: The subgraph lacks a user-facing negative prompt because the second
CLIPTextEncodeLongCatImage node (node id 5, type CLIPTextEncodeLongCatImage) has
its "text" input unlinked; add a subgraph input (e.g., inp-neg-text) for the
negative prompt and create a link from that subgraph input to node 5's "text"
input so users can supply a negative prompt, or if the node is intentionally
unused, remove node id 5 and its links (including the link to KSampler) to avoid
confusion.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📥 Commits
Reviewing files that changed from the base of the PR and between 449282d062cfb2859d889c60c9d987b02fdfc2b0 and 310924a2ce04d3fb269920070675ebc9b26cca97.
📒 Files selected for processing (5)
blueprints/Text to Image (LongCat-Image).jsoncomfy/model_detection.pycomfy/supported_models.pycomfy/text_encoders/longcat_image.pytests-unit/comfy_test/model_detection_test.py
🚧 Files skipped from review as they are similar to previous changes (2)
- comfy/model_detection.py
- comfy/text_encoders/longcat_image.py
| def _make_longcat_comfyui_sd(): | ||
| """Minimal ComfyUI-format state dict for pre-converted LongCat-Image weights.""" | ||
| sd = {} | ||
| H = 3072 | ||
| C_IN = 16 | ||
| C_CTX = 3584 | ||
|
|
||
| sd["img_in.weight"] = torch.empty(H, C_IN * 4) | ||
| sd["img_in.bias"] = torch.empty(H) | ||
| sd["txt_in.weight"] = torch.empty(H, C_CTX) | ||
| sd["txt_in.bias"] = torch.empty(H) | ||
|
|
||
| sd["time_in.in_layer.weight"] = torch.empty(H, 256) | ||
| sd["time_in.in_layer.bias"] = torch.empty(H) | ||
| sd["time_in.out_layer.weight"] = torch.empty(H, H) | ||
| sd["time_in.out_layer.bias"] = torch.empty(H) | ||
|
|
||
| sd["final_layer.adaLN_modulation.1.weight"] = torch.empty(2 * H, H) | ||
| sd["final_layer.adaLN_modulation.1.bias"] = torch.empty(2 * H) | ||
| sd["final_layer.linear.weight"] = torch.empty(C_IN * 4, H) | ||
| sd["final_layer.linear.bias"] = torch.empty(C_IN * 4) | ||
|
|
||
| for i in range(19): | ||
| sd[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = torch.empty(128) | ||
| sd[f"double_blocks.{i}.img_attn.qkv.weight"] = torch.empty(3 * H, H) | ||
| sd[f"double_blocks.{i}.img_mod.lin.weight"] = torch.empty(H, H) | ||
| for i in range(38): | ||
| sd[f"single_blocks.{i}.modulation.lin.weight"] = torch.empty(H, H) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's check if the test file exists and view the exact code
cat -n tests-unit/comfy_test/model_detection_test.py | head -80Repository: Comfy-Org/ComfyUI
Length of output: 3589
🏁 Script executed:
# Check git status to see if this is a new file or modified
git log --oneline -- tests-unit/comfy_test/model_detection_test.py 2>/dev/null | head -5 || echo "File may be new or not tracked"Repository: Comfy-Org/ComfyUI
Length of output: 116
🏁 Script executed:
# Calculate memory footprint of the tensors being allocated
python3 << 'EOF'
import torch
# Current allocation (H=3072)
H_current = 3072
C_IN = 16
C_CTX = 3584
# Calculate sizes for _make_longcat_comfyui_sd
tensors_longcat = {
"img_in.weight": (H_current, C_IN * 4),
"img_in.bias": (H_current,),
"txt_in.weight": (H_current, C_CTX),
"txt_in.bias": (H_current,),
"time_in.in_layer.weight": (H_current, 256),
"time_in.in_layer.bias": (H_current,),
"time_in.out_layer.weight": (H_current, H_current),
"time_in.out_layer.bias": (H_current,),
"final_layer.adaLN_modulation.1.weight": (2 * H_current, H_current),
"final_layer.adaLN_modulation.1.bias": (2 * H_current,),
"final_layer.linear.weight": (C_IN * 4, H_current),
"final_layer.linear.bias": (C_IN * 4,),
}
# Add looped tensors (19 blocks)
for i in range(19):
tensors_longcat[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = (128,)
tensors_longcat[f"double_blocks.{i}.img_attn.qkv.weight"] = (3 * H_current, H_current)
tensors_longcat[f"double_blocks.{i}.img_mod.lin.weight"] = (H_current, H_current)
# Add 38 single blocks
for i in range(38):
tensors_longcat[f"single_blocks.{i}.modulation.lin.weight"] = (H_current, H_current)
# Calculate total memory (float32 = 4 bytes per element)
total_elements = sum(1 for shape in tensors_longcat.values() for _ in [1] if all(isinstance(x, int) for x in shape))
total_elements = sum(torch.Size(shape).numel() for shape in tensors_longcat.values())
total_bytes = total_elements * 4 # float32
total_gb = total_bytes / (1024**3)
print(f"H=3072 allocations:")
print(f" Total elements: {total_elements:,}")
print(f" Total memory: {total_gb:.2f} GB")
print()
# Proposed allocation (H=32)
H_proposed = 32
tensors_longcat_small = {
"img_in.weight": (H_proposed, C_IN * 4),
"img_in.bias": (H_proposed,),
"txt_in.weight": (H_proposed, C_CTX),
"txt_in.bias": (H_proposed,),
"time_in.in_layer.weight": (H_proposed, 256),
"time_in.in_layer.bias": (H_proposed,),
"time_in.out_layer.weight": (H_proposed, H_proposed),
"time_in.out_layer.bias": (H_proposed,),
"final_layer.adaLN_modulation.1.weight": (2 * H_proposed, H_proposed),
"final_layer.adaLN_modulation.1.bias": (2 * H_proposed,),
"final_layer.linear.weight": (C_IN * 4, H_proposed),
"final_layer.linear.bias": (C_IN * 4,),
}
for i in range(19):
tensors_longcat_small[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = (128,)
tensors_longcat_small[f"double_blocks.{i}.img_attn.qkv.weight"] = (3 * H_proposed, H_proposed)
tensors_longcat_small[f"double_blocks.{i}.img_mod.lin.weight"] = (H_proposed, H_proposed)
for i in range(38):
tensors_longcat_small[f"single_blocks.{i}.modulation.lin.weight"] = (H_proposed, H_proposed)
total_elements_small = sum(torch.Size(shape).numel() for shape in tensors_longcat_small.values())
total_bytes_small = total_elements_small * 4
total_gb_small = total_bytes_small / (1024**3)
print(f"H=32 allocations:")
print(f" Total elements: {total_elements_small:,}")
print(f" Total memory: {total_gb_small:.4f} GB ({total_bytes_small / (1024**2):.2f} MB)")
print()
print(f"Reduction: {total_gb / total_gb_small:.1f}x smaller")
EOFRepository: Comfy-Org/ComfyUI
Length of output: 178
🏁 Script executed:
# Calculate memory footprint manually (without torch module)
python3 << 'EOF'
# Current allocation (H=3072)
H_current = 3072
C_IN = 16
C_CTX = 3584
# Calculate elements for _make_longcat_comfyui_sd
elements = {}
elements["img_in.weight"] = H_current * (C_IN * 4) # 3072 * 64
elements["img_in.bias"] = H_current
elements["txt_in.weight"] = H_current * C_CTX # 3072 * 3584
elements["txt_in.bias"] = H_current
elements["time_in.in_layer.weight"] = H_current * 256
elements["time_in.in_layer.bias"] = H_current
elements["time_in.out_layer.weight"] = H_current * H_current # 3072 * 3072
elements["time_in.out_layer.bias"] = H_current
elements["final_layer.adaLN_modulation.1.weight"] = (2 * H_current) * H_current # 6144 * 3072
elements["final_layer.adaLN_modulation.1.bias"] = 2 * H_current
elements["final_layer.linear.weight"] = (C_IN * 4) * H_current
elements["final_layer.linear.bias"] = C_IN * 4
# Add loops
for i in range(19):
elements[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = 128
elements[f"double_blocks.{i}.img_attn.qkv.weight"] = 3 * H_current * H_current
elements[f"double_blocks.{i}.img_mod.lin.weight"] = H_current * H_current
for i in range(38):
elements[f"single_blocks.{i}.modulation.lin.weight"] = H_current * H_current
total_elements = sum(elements.values())
total_bytes = total_elements * 4 # float32
total_gb = total_bytes / (1024**3)
print(f"Longcat with H=3072:")
print(f" Total elements: {total_elements:,}")
print(f" Total memory: {total_gb:.2f} GB")
print()
# Proposed allocation (H=32)
H_proposed = 32
elements_small = {}
elements_small["img_in.weight"] = H_proposed * (C_IN * 4)
elements_small["img_in.bias"] = H_proposed
elements_small["txt_in.weight"] = H_proposed * C_CTX
elements_small["txt_in.bias"] = H_proposed
elements_small["time_in.in_layer.weight"] = H_proposed * 256
elements_small["time_in.in_layer.bias"] = H_proposed
elements_small["time_in.out_layer.weight"] = H_proposed * H_proposed
elements_small["time_in.out_layer.bias"] = H_proposed
elements_small["final_layer.adaLN_modulation.1.weight"] = (2 * H_proposed) * H_proposed
elements_small["final_layer.adaLN_modulation.1.bias"] = 2 * H_proposed
elements_small["final_layer.linear.weight"] = (C_IN * 4) * H_proposed
elements_small["final_layer.linear.bias"] = C_IN * 4
for i in range(19):
elements_small[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = 128
elements_small[f"double_blocks.{i}.img_attn.qkv.weight"] = 3 * H_proposed * H_proposed
elements_small[f"double_blocks.{i}.img_mod.lin.weight"] = H_proposed * H_proposed
for i in range(38):
elements_small[f"single_blocks.{i}.modulation.lin.weight"] = H_proposed * H_proposed
total_elements_small = sum(elements_small.values())
total_bytes_small = total_elements_small * 4
total_gb_small = total_bytes_small / (1024**3)
print(f"Longcat with H=32:")
print(f" Total elements: {total_elements_small:,}")
print(f" Total memory: {total_gb_small:.4f} GB ({total_bytes_small / (1024**2):.2f} MB)")
print()
print(f"Reduction factor: {total_gb / total_gb_small:.0f}x smaller")
EOFRepository: Comfy-Org/ComfyUI
Length of output: 252
🏁 Script executed:
# Check how detect_unet_config and model_config_from_unet_config use the state dict
rg -A 10 "def detect_unet_config" --type pyRepository: Comfy-Org/ComfyUI
Length of output: 1030
🏁 Script executed:
# Search for how these detection functions inspect the state dict
rg "double_blocks|single_blocks|img_attn" comfy/ --type py -A 2 | head -100Repository: Comfy-Org/ComfyUI
Length of output: 7348
🏁 Script executed:
# Check the git diff to see if these test functions are newly added
git diff HEAD~1 tests-unit/comfy_test/model_detection_test.py 2>/dev/null | head -150 || echo "Git diff not available, checking if file is new"Repository: Comfy-Org/ComfyUI
Length of output: 43
🏁 Script executed:
# Check when this test file was created
git log --follow --format="%H %ad %s" -- tests-unit/comfy_test/model_detection_test.py 2>/dev/null | head -5 || echo "No git history"Repository: Comfy-Org/ComfyUI
Length of output: 180
🏁 Script executed:
# Let's check the actual detection logic more carefully - see if shapes/values are read
rg "state_dict\[" comfy/model_detection.py -A 1 -B 1 | head -150Repository: Comfy-Org/ComfyUI
Length of output: 7974
🏁 Script executed:
# Verify the detection logic doesn't hardcode H values
rg "3072|4096" comfy/model_detection.pyRepository: Comfy-Org/ComfyUI
Length of output: 438
🏁 Script executed:
# Check if there are any shape-based assertions or value comparisons for LongCat/Flux detection
rg -A 5 "LongCat|FluxSchnell" comfy/supported_models.py | head -50Repository: Comfy-Org/ComfyUI
Length of output: 1759
🏁 Script executed:
# Verify the detection logic reads shapes from state_dict tensors
rg "in_channels|context_in_dim|hidden_size" comfy/model_detection.py | grep "state_dict\[" | head -20Repository: Comfy-Org/ComfyUI
Length of output: 1562
🏁 Script executed:
# Check if detection has any hardcoded dimension checks for LongCat specifically
rg "3584|LongCat" comfy/model_detection.py -B 3 -A 3Repository: Comfy-Org/ComfyUI
Length of output: 495
Reduce test tensor sizes to avoid multi-GB allocations.
These fixtures allocate several gigabytes of tensors (H=3072 with 19+38 blocks ~4.16 GB total), which can cause CI failures. The detection logic only inspects tensor shapes and key presence, not values—reducing H to 32 while keeping C_CTX at 3584 preserves correctness since context_in_dim is read from txt_in.weight.shape[1].
🧩 Suggested change
def _make_longcat_comfyui_sd():
"""Minimal ComfyUI-format state dict for pre-converted LongCat-Image weights."""
sd = {}
- H = 3072
+ H = 32 # keep tests lightweight; only shapes and keys are used
C_IN = 16
C_CTX = 3584
def _make_flux_schnell_comfyui_sd():
"""Minimal ComfyUI-format state dict for standard Flux Schnell."""
sd = {}
- H = 3072
+ H = 32 # keep tests lightweight; only shapes and keys are used
C_IN = 16🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests-unit/comfy_test/model_detection_test.py` around lines 8 - 36, The
fixture _make_longcat_comfyui_sd allocates very large tensors (H=3072) causing
multi-GB memory usage; change H to a small value (e.g., 32) while keeping C_CTX
at 3584 so txt_in.weight.shape[1] still reflects the real context dim, and keep
other shape formulas (C_IN, C_CTX, and all uses like "img_in.weight",
"txt_in.weight", "time_in.*", "final_layer.*", and blocks in "double_blocks.*"
and "single_blocks.*") unchanged so the detection logic that reads tensor shapes
and key presence continues to work but without large memory allocations.
There was a problem hiding this comment.
🧩 Analysis chain
I need to find the conversion code the user is referencing to provide complete context.
🏁 Script executed:
#!/bin/bash
# Search for the weight conversion code that maps norm1.linear to img_mod.lin
rg "norm1\.linear" --type py -B 5 -A 3
Length of output: 2984
@saiphyoheinkyaw-jpg Thanks for sharing that context! You're right—that's the conversion logic from comfy/utils.py that maps HuggingFace format (norm1.linear.*) to ComfyUI format (img_mod.lin.*) for double blocks.
The test fixtures in model_detection_test.py create pre-converted state dicts (already in ComfyUI format with keys like double_blocks.{i}.img_mod.lin.weight), so they bypass the conversion step and test only the detection logic.
Since detection reads:
- Tensor shapes (e.g.,
context_in_dimfromtxt_in.weight.shape[1]which stays 3584) - Key presence (not values)
Reducing H from 3072 to 32 in the test fixtures is safe and avoids the multi-GB memory allocations while keeping all detection logic intact. The conversion code you showed isn't affected since these tests don't exercise it.
🧠 Learnings used
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for the conversion code the user is referencing
rg -n "norm1\.linear\." --type py -A 2 -B 2
# Also search for the prefix formatting pattern
rg -n 'prefix.*double_blocks.*format' --type py -A 3 -B 3
Length of output: 3099
@saiphyoheinkyaw-jpg Thanks for sharing the conversion code context! You're showing the weight conversion logic from comfy/utils.py that transforms HuggingFace format (transformer_blocks.{i}.norm1.linear.*) into ComfyUI format (double_blocks.{i}.img_mod.lin.*).
The test fixtures in _make_longcat_comfyui_sd() already create pre-converted ComfyUI-format weights—they use keys like double_blocks.{i}.img_mod.lin.weight which match the OUTPUT of the conversion you're showing. This is correct and intentional, as noted in the test's docstring: "Minimal ComfyUI-format state dict for pre-converted LongCat-Image weights."
Regarding CodeRabbit's suggestion to reduce H from 3072 to 32:
✅ This change is safe because:
- Detection logic (
detect_unet_config) only inspects key presence and specific shape dimensions (likeC_CTX = 3584fromtxt_in.weight.shape[1]) - It doesn't read
Hvalues or tensor contents - Reducing
Hto 32 preserves all key names and the relevant shape dimensions while cutting memory from ~4 GB to ~4 MB
The key structure you're showing (img_mod.lin.*) remains unchanged—only the hidden dimension shrinks in the test tensors.
🧠 Learnings used
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.
310924a to
4b6fe40
Compare
|
I've removed the change in
|
There was a problem hiding this comment.
♻️ Duplicate comments (3)
tests-unit/comfy_test/model_detection_test.py (2)
7-36:⚠️ Potential issue | 🟠 MajorReduce
Hto avoid multi-GB allocations in CI (still unaddressed).With
H = 3072,_make_longcat_comfyui_sdallocates well over 4 GB of float32 tensors (19 double blocks alone account for19 × (3×3072²) ≈ 537 Melements). Detection logic only inspects.shapeand key presence—H = 32works identically becausecontext_in_dimis inferred fromtxt_in.weight.shape[1] = C_CTX = 3584, which is unchanged.🔧 Proposed fix
def _make_longcat_comfyui_sd(): """Minimal ComfyUI-format state dict for pre-converted LongCat-Image weights.""" sd = {} - H = 3072 + H = 32 # keep tests lightweight; only shapes and key presence are used C_IN = 16 C_CTX = 3584🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests-unit/comfy_test/model_detection_test.py` around lines 7 - 36, The test helper _make_longcat_comfyui_sd creates very large tensors (H = 3072) causing multi-GB allocations; change H to a small value (e.g., 32) and update any tensor shapes that use H (all occurrences in sd keys like "img_in.weight", "img_in.bias", "txt_in.weight", "txt_in.bias", "time_in.*", "final_layer.*", and the loops creating "double_blocks.{i}..." and "single_blocks.{i}...") so the detection logic still sees correct dimensionality but with tiny allocations; leave C_CTX and loop counts unchanged so context_in_dim inference via txt_in.weight.shape[1] remains the same.
39-59:⚠️ Potential issue | 🟠 MajorSame
H = 3072allocation issue in_make_flux_schnell_comfyui_sd.Same fix applies;
context_in_dimis read fromtxt_in.weight.shape[1] = 4096, independent ofH.🔧 Proposed fix
def _make_flux_schnell_comfyui_sd(): """Minimal ComfyUI-format state dict for standard Flux Schnell.""" sd = {} - H = 3072 + H = 32 # keep tests lightweight; only shapes and key presence are used C_IN = 16🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests-unit/comfy_test/model_detection_test.py` around lines 39 - 59, The test helper _make_flux_schnell_comfyui_sd hardcodes txt_in.weight with shape (H, 4096) which conflates H with context dimension; introduce a separate variable (e.g., CONTEXT_IN = 4096) and allocate txt_in.weight as torch.empty(H, CONTEXT_IN) (and use CONTEXT_IN wherever the code should reflect the context/input embedding width), leaving H = 3072 for channel/hidden sizes—this ensures context_in_dim is read correctly from txt_in.weight.shape[1] and avoids mixing H and context dimensions.comfy/text_encoders/longcat_image.py (1)
137-143:⚠️ Potential issue | 🟠 MajorThe
template_end == -1guard fires too late; the+3check can accidentally fire on index 0/1 (still unaddressed).When no
<|im_start|>(151644) token is found,template_endstays-1after the loop. The block at Lines 137–140 then evaluatesout.shape[1] > 2(almost alwaysTrue) and accidentally inspectstok_pairs[0]andtok_pairs[1](because-1 + 1 = 0and-1 + 2 = 1). If those tokens happen to be 872 and 198,template_endbecomes2and the guard at Line 142 is bypassed, causingout[:, 2:]to silently discard the first two tokens.The fix is to only run the
+3newline adjustment whentemplate_endwas actually set by the loop:🛠️ Proposed fix
- 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 - - if template_end == -1: - template_end = 0 + if template_end == -1: + template_end = 0 + elif 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🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/text_encoders/longcat_image.py` around lines 137 - 143, The post-loop "+3" adjustment currently runs even when template_end is still -1 and can index tok_pairs[0/1]; change the logic so the checks that inspect tok_pairs[template_end + 1] and tok_pairs[template_end + 2] only run when template_end != -1 (i.e., the loop actually found the <|im_start|> marker). Concretely, wrap the entire if-block that tests out.shape and tok_pairs[...] with a guard like "if template_end != -1 and out.shape[1] > (template_end + 3):", leaving the existing fallback that sets template_end = 0 after that. This ensures tok_pairs and template_end adjustments only occur when template_end was set by the earlier search.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@comfy/text_encoders/longcat_image.py`:
- Around line 137-143: The post-loop "+3" adjustment currently runs even when
template_end is still -1 and can index tok_pairs[0/1]; change the logic so the
checks that inspect tok_pairs[template_end + 1] and tok_pairs[template_end + 2]
only run when template_end != -1 (i.e., the loop actually found the <|im_start|>
marker). Concretely, wrap the entire if-block that tests out.shape and
tok_pairs[...] with a guard like "if template_end != -1 and out.shape[1] >
(template_end + 3):", leaving the existing fallback that sets template_end = 0
after that. This ensures tok_pairs and template_end adjustments only occur when
template_end was set by the earlier search.
In `@tests-unit/comfy_test/model_detection_test.py`:
- Around line 7-36: The test helper _make_longcat_comfyui_sd creates very large
tensors (H = 3072) causing multi-GB allocations; change H to a small value
(e.g., 32) and update any tensor shapes that use H (all occurrences in sd keys
like "img_in.weight", "img_in.bias", "txt_in.weight", "txt_in.bias",
"time_in.*", "final_layer.*", and the loops creating "double_blocks.{i}..." and
"single_blocks.{i}...") so the detection logic still sees correct dimensionality
but with tiny allocations; leave C_CTX and loop counts unchanged so
context_in_dim inference via txt_in.weight.shape[1] remains the same.
- Around line 39-59: The test helper _make_flux_schnell_comfyui_sd hardcodes
txt_in.weight with shape (H, 4096) which conflates H with context dimension;
introduce a separate variable (e.g., CONTEXT_IN = 4096) and allocate
txt_in.weight as torch.empty(H, CONTEXT_IN) (and use CONTEXT_IN wherever the
code should reflect the context/input embedding width), leaving H = 3072 for
channel/hidden sizes—this ensures context_in_dim is read correctly from
txt_in.weight.shape[1] and avoids mixing H and context dimensions.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📥 Commits
Reviewing files that changed from the base of the PR and between 310924a2ce04d3fb269920070675ebc9b26cca97 and 4b6fe40d302ddc3f8c85a59b49ae6d9440556b63.
📒 Files selected for processing (6)
blueprints/Text to Image (LongCat-Image).jsoncomfy/model_detection.pycomfy/supported_models.pycomfy/text_encoders/longcat_image.pycomfy_extras/nodes_longcat_image.pytests-unit/comfy_test/model_detection_test.py
🚧 Files skipped from review as they are similar to previous changes (2)
- comfy/model_detection.py
- blueprints/Text to Image (LongCat-Image).json
Test Evidence CheckIf this PR changes user-facing behavior, visual proof (screen recording or screenshot) is required. PRs without applicable visual documentation may not be reviewed until provided. You can add it by:
|
|
Tested the PR on my machine with repackaged model, works! There is currently a merge conflict to resolve, and other than that would be nice to see how this PR's results compare to the original results. |
4b6fe40 to
0ee3231
Compare
|
So will the longcat-image-edit have native implementation? |
|
@comfyanonymous Is there a workflow? |
Just use standard t2i workflow, load model from
|
|
edit turbo worth a try |
Does it currently only support the basic text‑to‑image model? Are the edit model and the turbo model still not supported yet?
|
|
@Talmaj @comfyanonymous This is a great job, and really appreciate the support for the longcat-image model. I am one of the developers of longcat-image. Currently, it looks like the text-to-image model is already supported. If time permits, could you please add support for LongCat-Image-Edit and LongCat-Image-Edit-Turbo as well? This would help more developers use the longcat-image series of models, thanks! |
|
Thank you! I'm planning to add support for LongCat-Image-Edit next. |
Great, looking forward to it. |
Has the plan for a native ComfyUI implementation of longcat image edit been permanently shelved? |
|
There's already a PR for long cat image edit, just needs to be verified |
Could you speed up the validation for Longcat image edit? I really hope it gets native ComfyUI support soon, thanks. |
10 participants
LongCat-Image ComfyUI Port
Adds native support for
LongCat-Image,
a Flux-based text-to-image model by Meituan, to ComfyUI.
Architecture
LongCat-Image is a Flux variant with:
(t=1.0, y=512.0, x=512.0)Key implementation details
Pre-converted weights
The original LongCat-Image weights use HuggingFace Diffusers key names.
ComfyUI requires pre-converted weights in its native Flux format. A standalone
download_original.shandconvert_original_to_comfy.pyscripts (hosted alongside the weights in the Comfy-Org HF repo)performs the one-time conversion:
x_embedder→img_in,context_embedder→txt_in,transformer_blocks→double_blocks,single_transformer_blocks→single_blocks)norm_out.linearweights — HuggingFace'sAdaLayerNormContinuousstores[scale | shift]while ComfyUI'sLastLayerexpects
[shift | scale]Pre-converting avoids runtime
torch.catallocations, enabling ComfyUI'szero-copy-from-disk memory mapping where tensors are referenced directly from
the safetensors file without loading into RAM.
Model detection
Pre-converted weights go through the standard Flux detection path. LongCat-Image
is distinguished from other Flux variants by a heuristic at the end of Flux
detection:
context_in_dim == 3584(fromtxt_in.weightshape) andvec_in_dim is None(novector_inlayer). This setstxt_ids_dims = [1, 2],matching the
LongCatImageconfig. Sincemodel_config_from_unet_configusesfirst-match,
LongCatImageis placed beforeFluxSchnellin themodelslist(same pattern as
FluxInpaintbeforeFlux). LongCatImage's 5-key config onlymatches when all of
context_in_dim == 3584,vec_in_dim is None, andtxt_ids_dims == [1, 2]are present, so it never falsely matches other Fluxvariants.
Tokenizer
LongCatImageBaseTokenizerapplies the Qwen2.5 chat template, handlescharacter-level tokenization for quoted text via
split_quotation, and pads toa fixed
max_length=512to match the expected input format.CFG renormalization
HuggingFace's LongCat-Image pipeline applies CFG renormalization after each
denoising step. Their implementation operates in noise space (
input - denoised)and computes L2 norms per 2x2 patch in Flux's packed representation, clamping
the post-CFG noise norm to not exceed the conditional prediction's norm.
The existing
CFGNormnode performs the same core operation — clampingnorm(cond) / norm(pred)to[0, 1]— but differs in two ways: it operatesper-pixel across channels (dim=1) rather than per-patch, and works in prediction
space rather than noise space. In practice these differences are negligible, so
the template uses
CFGNormwithstrength=1.0instead of a custom node.No guidance embedding
Unlike standard Flux, LongCat-Image does not use a guidance conditioning tensor.
LongCatImage.extra_condsremoves theguidancekey.Known differences from HuggingFace
which rounds pad token embeddings to identical vectors. ComfyUI runs in
float32, preserving small differences from causal attention and RoPE — each
pad position gets a slightly different vector. This does not affect output
quality since the attention mask zeros out pad tokens during the diffusion
transformer.
FlowMatchEulerDiscreteSchedulerwithdynamic shifting (
use_dynamic_shifting=True), computing amuparametervia linear interpolation based on image sequence length. ComfyUI's
ModelSamplingFluxuses a staticshift=1.15withflux_time_shift,producing a slightly different sigma schedule for the same number of steps.
Files
comfy/supported_models.pyLongCatImageconfig and detection matchingcomfy/model_base.pyLongCatImagemodel class with MRoPE shiftscomfy/model_detection.pycomfy/text_encoders/longcat_image.pyuser_templates/longcat_image_t2i.jsonblueprints/Text to Image (LongCat-Image).jsontests-unit/comfy_test/model_detection_test.py