[model] Flux2 Klein Port - #1349
Conversation
mergify
Bot
added
scope: inference
|
Merge ProtectionsYour pull request matches the following merge protections and will not be merged until they are valid. 🔴 PR merge requirementsWaiting for
This rule is failing.
|
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements support for the Flux2 model family, focusing on the Flux2 Klein distilled variant. It introduces the Flux2 transformer, Qwen3 text encoder, and Flux2 VAE, along with specialized pipeline stages for packed latent handling and resolution-dependent timestep shifting. Review feedback highlights several issues: missing **kwargs in Flux2Attention.forward which may cause type errors, unused lora_scale variables that break LoRA support, and leftover debugging logic in transformer blocks. The implementation also currently lacks support for multiple text encoders and assumes square images in its RoPE fallback logic. A potential compatibility issue with the enable_thinking parameter in apply_chat_template was also noted.
| def forward( | ||
| self, | ||
| hidden_states: torch.Tensor, | ||
| encoder_hidden_states: Optional[torch.Tensor] = None, | ||
| freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, | ||
| ) -> torch.Tensor: |
There was a problem hiding this comment.
The forward method of Flux2Attention is missing **kwargs in its signature. This will cause a TypeError if joint_attention_kwargs contains any keys other than the ones explicitly filtered out in Flux2TransformerBlock.forward (e.g., LoRA parameters or other custom attention attributes). Since this class inherits from AttentionModuleMixin, it should ideally follow the standard signature that accepts extra arguments.
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs,
) -> torch.Tensor:| lora_scale = joint_attention_kwargs.pop("scale", 1.0) | ||
| else: | ||
| lora_scale = 1.0 |
There was a problem hiding this comment.
The lora_scale is extracted from joint_attention_kwargs but never used in the forward pass. This effectively ignores any LoRA scaling factors passed through the pipeline, which will lead to incorrect results when using LoRA weights. If LoRA support is intended, this scale should be applied to the relevant layers or passed down to them.
| debug_enc = joint_attention_kwargs.get("_debug_double_enc") | ||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["context_attn_output"].append(context_attn_output.detach().clone()) | ||
| encoder_hidden_states = encoder_hidden_states + context_attn_output | ||
|
|
||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["after_attn"].append(encoder_hidden_states.detach().clone()) | ||
|
|
||
| norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) | ||
| norm_encoder_hidden_states = ( | ||
| norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp | ||
| ) | ||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["before_ff"].append(norm_encoder_hidden_states.detach().clone()) | ||
|
|
||
| context_ff_output = self.ff_context(norm_encoder_hidden_states) | ||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["context_ff_output"].append(context_ff_output.detach().clone()) | ||
| debug_enc["c_gate_mlp"].append(c_gate_mlp.detach().clone()) | ||
| context_ff_update = c_gate_mlp * context_ff_output | ||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["context_ff_update"].append(context_ff_update.detach().clone()) | ||
| encoder_hidden_states = encoder_hidden_states + context_ff_update | ||
| if encoder_hidden_states.dtype == torch.float16: | ||
| encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) | ||
|
|
||
| if debug_enc is not None and joint_attention_kwargs.get("_double_block_index") == debug_enc.get("block_index"): | ||
| debug_enc["after_ff"].append(encoder_hidden_states.detach().clone()) | ||
|
|
There was a problem hiding this comment.
The Flux2TransformerBlock.forward method contains intrusive debugging logic (using _debug_double_enc and _double_block_index) that detaches and clones tensors into a dictionary. This logic appears to be leftover from parity testing and adds unnecessary overhead and clutter to the production model code. It should be removed or moved to a dedicated debugging utility.
| if isinstance(encoder_hidden_states, (list, tuple)): | ||
| encoder_hidden_states = encoder_hidden_states[0] |
There was a problem hiding this comment.
The model is hardcoded to use only the first text encoder output if a list is provided. While this works for Flux2 Klein (which uses a single Qwen3 encoder), it will break or produce incorrect results for the full Flux model (which uses two encoders: CLIP and T5), despite this class being registered as FluxTransformer2DModel. Consider adding support for multiple encoders or at least a warning/TODO for future-proofing.
| img_h, img_w = h, w | ||
| else: | ||
| img_seq_len = hidden_states.shape[1] | ||
| img_h = img_w = int(img_seq_len ** 0.5) |
There was a problem hiding this comment.
The fallback logic for inferring image dimensions from img_seq_len assumes square images (int(img_seq_len ** 0.5)). This will produce incorrect RoPE positional embeddings for rectangular images if the input is passed as a 3D tensor (not 5D). While FastVideo typically uses 5D tensors, the model should be robust to non-square aspect ratios in all input formats it claims to support.
| formatted = tokenizer.apply_chat_template( | ||
| messages, | ||
| tokenize=False, | ||
| add_generation_prompt=True, | ||
| enable_thinking=False, | ||
| ) |
There was a problem hiding this comment.
The call to tokenizer.apply_chat_template includes enable_thinking=False. This parameter was introduced in very recent versions of the transformers library and may not be supported by all tokenizers or older versions of the library. To ensure better compatibility across different environments, consider removing this argument or wrapping it in a check.
formatted = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
- L40S TP2/TP4 has bounded BF16 TP drift. - H100 image artifacts were generated at 1024x1024 and look good. - Full model image conditioning/caption upsampling remains out of scope.
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @SolitaryThinker, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
|
/merge |
…ehavior (hao-ai-lab#1349) Shared pipeline/loader paths touched by the Flux2 port now no-op for existing models, with all behavior changes gated behind Flux2-only conditions: - text_encoding: restore the original chat-template call (tok.apply_chat_template(processed_texts, **tok_kwargs)) for pre-formatted chat models (HunyuanVideo 1.5 / Qwen2.5-VL). This recovers the dropped add_generation_prompt and the inner tokenizer; the new two-step formatting now applies only to Flux2 Klein's raw-string path. - component_loader: gate the processor_config.json AutoProcessor shortcut behind a new require_processor flag (set only on Mistral3), so existing encoders stay on AutoTokenizer even when their tokenizer dir ships a processor_config.json. - decoding: gate the 5D->4D squeeze on _is_flux2_packed so video VAEs decoding T=1 latents are untouched. - denoising: revert incidental non-Flux image_embeds / prompt_embeds changes back to the original. - fsdp_load / decoding: document the existing Flux2-only gates.
…oke port main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). It returns BaseEncoderOutput(last_hidden_state, hidden_states), supports GQA + the same qkv/gate_up fusion, and update_model_arch populates Z-Image-Turbo's dims (2048/24/16) from its config.json — so our parallel Qwen3Model was ~540 lines of duplicate. Drop the bespoke encoder + config (resolved during the rebase onto main; the files are now identical to main) and wire Z-Image to the shared encoder: - registry: map Z-Image-Turbo's "Qwen3Model" architecture string -> Qwen3ForCausalLM. - encoder parity test: use Qwen3TextConfig; move the strict-load allowlist ({"lm_head.weight"}) from a class attribute to a test-side constant, since the shared encoder's loader is intentionally lenient (serves multiple models). - PORT_STATUS / README: mark the encoder "reused"; the prior A40 parity PASS was measured against the removed bespoke encoder, so text-encoder parity is flagged RE-VALIDATION PENDING against the shared class. Scheduler/tokenizer/VAE parity are unaffected and still PASS. The gated scheduler option (use_reference_discrete_timesteps, default False) is unchanged and still needed for Z-Image timestep parity.
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
…positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (hao-ai-lab#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).
3 participants
Summary
Re-opening of #1345 from a fork I have push access to. The original branch on @Gnav3852's fork was not push-accessible from this environment after the review fixes were applied.
All 12 original commits remain authored by @Gnav3852 (verified via `jj metaedit --author` per-commit). Five additional fix commits by @SolitaryThinker address the review findings on #1345:
Plus a `[docs]:` commit updating `tests/local_tests/flux2/README.md`.
Review findings addressed
From the original review of #1345:
Deferred (not in this PR)
Provenance
Test plan