Twinflow-Z-Image Turbo support#13285
Open
azazeal04 wants to merge 17 commits intoComfy-Org:masterfrom
Open
Conversation
Add compatibility for TwinFlow-Z-Image LoRAs to target t_embedder_2.* keys.
Added TwinFlow_Z_Image class for new model integration.
Add twinflow_z_image_key_mapping function for key mapping.
Added a method to compute TwinFlow adaLN input with delta-time conditioning.
azazeal04
requested review from
Kosinkadink,
comfyanonymous and
guill
as code owners
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: de2ff57f3c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
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 a TwinFlow‑Z‑Image diffusion transformer at 🚥 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: 7
🤖 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/ldm/twinflow/model.py`:
- Around line 627-648: The pad logic extends cap_feats and cap_pos_ids but
forgets to pad cap_mask, causing length mismatches; inside the if
self.pad_tokens_multiple block (after cap_pos_ids is padded) compute pad_extra
as already done and if pad_extra > 0 extend cap_mask to match cap_feats by
concatenating a pad tensor of shape (cap_mask.shape[0], pad_extra) on the
sequence dimension using cap_mask.dtype and cap_mask.device, with the pad value
set to the same "invalid" value used by incoming masks (0/False); ensure this
updated cap_mask is what gets passed into context_refiner's layer(...) calls.
- Around line 529-537: transformer_options["target_timestep"] may be a plain
int/float, but the code assumes it's a tensor and calls .to(...) on it (see
transformer_options, target_timestep, target_t, and usage in the t embedding
block), which raises AttributeError for scalars; fix by converting
target_timestep to a tensor robustly with torch.as_tensor(target_timestep,
device=t.device, dtype=t.dtype) (then use target_t = that tensor and
expand_if_needed before using self.t_embedder_2 and combining with t_emb) so
both scalar and tensor inputs are accepted.
In `@comfy/lora_convert.py`:
- Around line 35-41: The current twinflow_z_image_lora_to_diffusers function can
silently overwrite existing "t_embedder_2.*" entries when remapping
"t_embedder.*" keys; change the move logic to first check if new_key already
exists in state_dict and if so do not overwrite it — either skip moving the
original key (preserve the existing t_embedder_2 tensor) and optionally log or
raise a descriptive warning/error; update the code path that does
state_dict[new_key] = state_dict.pop(key) to perform an existence check for
new_key and handle collisions without replacing existing tensors.
In `@comfy/lora.py`:
- Around line 331-337: The aliasing loop in comfy/lora.py for
comfy.model_base.TwinFlow_Z_Image blindly adds "t_embedder_2.*" entries from
"t_embedder.*" and can overwrite existing native "t_embedder_2" mappings; change
the loop that iterates over list(key_map.keys()) so that before assigning
key_map[key.replace("t_embedder.", "t_embedder_2.", 1)] = key_map[key] you first
check whether the replacement key already exists and only add the alias when it
does not (i.e., skip or preserve existing "t_embedder_2.*" entries), leaving all
other logic and symbol names (model, key_map, TwinFlow_Z_Image) unchanged.
In `@comfy/model_detection.py`:
- Around line 47-79: The TwinFlow_Z_Image branch returns dit_config before
applying the Z-Image allow_fp16 heuristic, so add an allow_fp16 entry to
dit_config (e.g., dit_config["allow_fp16"]=True/False) using the same heuristic
used by ZImage.__init__ (inspect the existing Z-Image logic that checks
dimensionality/axes or cap_embedder shapes to decide fp16 eligibility) and set
it right before the return in the TwinFlow_Z_Image detection block (the block
that builds dit_config for "TwinFlow_Z_Image" / "twinflow_z_image"). Ensure you
reference the same signals (dit_config["dim"], dit_config["axes_lens"] or
cap_embedder_key-derived shapes, and pad token checks) so TwinFlow_Z_Image
mirrors Z-Image's fp16 gating.
In `@comfy/utils.py`:
- Around line 821-829: The current twinflow_z_image_key_mapping function moves
tensors with keys starting "t_embedder_2." using state_dict.pop which can
overwrite existing "t_embedder." entries and deletes the original; instead
perform a non-destructive remap: for each key starting with "t_embedder_2.",
compute new_key = key.replace("t_embedder_2.", "t_embedder.", 1) and copy the
tensor (state_dict[new_key] = state_dict[key]) without popping; if new_key
already exists, leave the existing state_dict[new_key] intact and do not delete
state_dict[key] (i.e., skip overwrite), ensuring both t_embedder_2.* and
t_embedder.* remain present to preserve backward compatibility in
twinflow_z_image_key_mapping.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 18910f50-263d-4ae9-b713-c2f3f85ec166
📒 Files selected for processing (8)
comfy/ldm/twinflow/model.pycomfy/lora.pycomfy/lora_convert.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/utils.py
Ensure compatibility by mapping keys for t_embedder.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/ldm/twinflow/model.py`:
- Around line 110-113: The hot path currently does unnecessary
unsqueeze/repeat/flatten when n_rep == 1; change the condition that checks n_rep
(computed from n_local_heads and n_local_kv_heads) to only perform the
unsqueeze/repeat/flatten on xk and xv when n_rep > 1, and otherwise leave xk and
xv unchanged so no extra allocations/copies occur; update the block around
n_rep, xk, xv, n_local_heads, and n_local_kv_heads to use if n_rep > 1 and keep
downstream shape expectations intact when skipping the replication.
- Line 708: The TwinFlow adaLN input computed by _compute_twinflow_adaln
(adaln_input, which encodes t_embedder_2/target_timestep dual-timestep
conditioning) is being overwritten/rebuilt from t_emb when pooled CLIP is
enabled; preserve and propagate adaln_input instead of replacing it with a
t_emb-only reconstruction. Update the block that currently rebuilds conditioning
from t_emb (the code paths around adaln_input and t_emb handling when pooled
CLIP is active) so that pooled CLIP uses the already-computed adaln_input (and
retains t_embedder_2/target_timestep signals) rather than discarding them, and
remove the t_emb-only rebuild logic so backward-compatible dual-timestep
conditioning remains intact.
- Around line 696-703: The _forward function uses a shared mutable default
transformer_options={}, which is mutated (keys "total_blocks", "block_type",
"block_index") and can leak state across calls; change the signature default to
transformer_options=None and inside _forward create a fresh dict (e.g., opts =
dict(transformer_options or {})) then mutate and use opts instead of
transformer_options so each invocation has its own options object and avoids
concurrency/state issues.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d385ed21-a0cd-478d-81cd-9a1e3c1ec1a8
📒 Files selected for processing (1)
comfy/ldm/twinflow/model.py
Ensure new key is added only if it doesn't exist.
Prevent overwriting existing keys when mapping t_embedder_2 keys to t_embedder.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/ldm/twinflow/model.py`:
- Around line 679-684: The code in patchify_and_embed clears the caption mask by
setting mask = None which breaks variable-length caption handling for the joint
transformer; instead construct and return a combined attention mask (preserving
cap_mask) that covers caption tokens and image patches so self.layers and
self.context_refiner see the correct masking. Update the return from
patchify_and_embed to compute mask from cap_mask (and any patch masks if
applicable), account for pad_tokens_multiple padding and batch sizes, and return
that mask alongside padded_full_embed, img_sizes, l_effective_cap_len, freqs_cis
so downstream callers (e.g., self.layers and self.context_refiner) receive the
proper attention_mask.
- Around line 525-527: The pad token parameters x_pad_token and cap_pad_token
are created with torch.empty(), leaving them uninitialized if a checkpoint
conversion doesn't populate both; change their initialization to a deterministic
fallback (e.g., zero-filled tensor) when pad_tokens_multiple is not None so
downstream runs don't read garbage. Locate the branch that checks
pad_tokens_multiple and replace the torch.empty-based init of nn.Parameter for
x_pad_token and cap_pad_token with a zero-initialized tensor (use torch.zeros
with the same shape, device, and dtype) so loaded weights can still overwrite
them safely.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 5f48ee12-d276-4b23-be3e-774efd853ec5
📒 Files selected for processing (1)
comfy/ldm/twinflow/model.py
Add logic to set allow_fp16 based on weight standard deviation
Author
|
Model used: https://huggingface.co/azazeal2/TwinFlow-Z-Image-Turbo-repacked/tree/main/ComfyUI |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added support for the Twinflow Z-Image model