Add FLUX.2-klein-base-9B contrib model#146
Conversation
Add 2048x2048 benchmark results (191.44s, TP=4, 50 steps) validated by xniwang on trn2.3xlarge. Document that 4096x4096 is NOT SUPPORTED due to a fundamental model limitation (max_area=4MP), not a hardware constraint. The model produces noise/gray at 4K on ALL devices including H100. Maximum supported resolution: 2048x2048.
| for _ in range(2): | ||
| app( | ||
| prompt=args.prompt, | ||
| negative_prompt=args.negative_prompt, |
There was a problem hiding this comment.
Flux2KleinPipeline doesn't seem to accept negative_prompt.
|
|
||
| os.environ["LOCAL_WORLD_SIZE"] = str(self.config.neuron_config.world_size) | ||
| if _HARDWARE == hardware.TRN2: | ||
| os.environ["NEURON_RT_VIRTUAL_CORE_SIZE"] = "2" |
There was a problem hiding this comment.
I feel like it's a bad idea to set a runtime variable at compile time, in a function like get_compiler_args. Above, you decide which kernel to use based on if it's sharded or not, but it makes that decision based on this environment variable that's presumably set at compile time.
Philosophically this shouldn't happen because you could theoretically load multiple (pre-compiled) NEFFs in a single runtime.
| # For FLUX.2-klein: latent is (B, H*W, 128) after pack | ||
| # At 1024x1024: H=W=64 (1024/16), so img_seq = 4096 | ||
| vae_scale_factor = getattr(self.config, "vae_scale_factor", 16) | ||
| num_patches = self.config.height * self.config.width // (vae_scale_factor**2) |
There was a problem hiding this comment.
This assumes width is perfectly divisible by vae_scale_factor**2, which could result in the input being truncated. I can see why this is the way, but an assertion or a warning if this occurs might be helpful.
|
Approved because I managed to generate an image using the same prompt. Comments left are non-blocking. |
2 participants
Note: The below template includes items meant for model contributions only. For other contributions such as bug fixes, features, etc., only fill out the relevant portions of the form.
Description
Add FLUX.2-klein-base-9B (9.08B parameter diffusion transformer) as a contrib model with NxD Inference tensor parallelism on trn2.3xlarge.
FLUX.2-klein differs from FLUX.1 in several key ways: SwiGLU activation (vs GELU), pre-computed modulation, fused QKV+MLP projections in single-stream blocks, Qwen3-8B text encoder, and 32 latent channels with 4D RoPE. The implementation splits all fused SwiGLU projections into separate ColumnParallelLinear layers for correct TP sharding, and decomposes the massive fused
to_qkv_mlp_projinto independent Q/K/V and MLP projections to stay within compiler instruction limits.Model Information
Model Name: FLUX.2-klein-base-9B
Model Architecture: Diffusion transformer (DiT) with 8 double-stream MMDiT blocks + 24 single-stream DiT blocks, 4D RoPE, SwiGLU, pre-computed modulation
Purpose: Text-to-image generation
Checklist
Required Components
Accuracy Test (
test/integration/test_model.py)README.md with the following sections:
Source Code (
src/)modeling_flux2_klein.py: Full NxDI model implementation (~1350 lines)application.py: Config factory, NeuronTransformerWrapper, NeuronFlux2KleinApplicationgenerate_flux2_klein.py: CLI entry pointOptional Components
Folder Structure
Testing
How did you test this change?
Tested on trn2.3xlarge (LNC=2, TP=4) with Neuron SDK 2.29 (DLAMI 20260410). Compiled backbone, ran direct backbone comparison against HF CPU reference with identical inputs, then ran full end-to-end pipeline with 5 warm generations for benchmarking.
Test Results:
Backbone accuracy:
Benchmark (1024x1024, 30 steps, guidance_scale=4.0, classic CFG):
Compatibility
Tested with:
Additional Information