Flux2klein Onboarding#445
Conversation
amepas
force-pushed
the
flux2klein-onboarding-clean
branch
from
e4e7f55 to
8a38bd9
Compare
amepas
force-pushed
the
flux2klein-onboarding-clean
branch
2 times, most recently
from
b22d0ce to
a67c8b7
Compare
amepas
force-pushed
the
flux2klein-onboarding-clean
branch
2 times, most recently
from
48b898f to
20a03d2
Compare
amepas
force-pushed
the
flux2klein-onboarding-clean
branch
4 times, most recently
from
7bcb105 to
f446205
Compare
amepas
force-pushed
the
flux2klein-onboarding-clean
branch
from
f446205 to
56e6a4d
Compare
There was a problem hiding this comment.
I still see a few print statement instead of max_logging, but I think it's fine for now, we can send out a small PR for that later. RN let's prioritize merging the model onboarding PR.
I was able to reproduce the same e2e generation time as reported.
|
🤖 Hi @eltsai, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @eltsai, but I was unable to process your request. Please see the logs for more details. |
|
Seems like the linter test failed because of #443 , but #443 itself passed the linter tests, @syhuang22 is investigating it now |
| )["params"] | ||
|
|
||
|
|
||
| class FlaxSwiGluFeedForward(nn.Module): |
There was a problem hiding this comment.
New models in maxdiffusion should be using nnx over linen. Is there a reason why linen was used?
| interactive: False | ||
|
|
||
| # 9B Architecture Dimensions | ||
| depth: 24 # num_single_layers |
There was a problem hiding this comment.
why are these required? They can be retrieved from the huggingface config.
| import jax.numpy as jnp | ||
| import jax | ||
|
|
||
| for k, v in list(d.items()): |
There was a problem hiding this comment.
should filter out norm layers to keep those calculations in fp32 for numerical stability. For example, see https://github.com/AI-Hypercomputer/maxdiffusion/blob/main/src/maxdiffusion/pipelines/wan/wan_pipeline.py#L95
| first_leaf = jax.tree_util.tree_leaves(params)[0] | ||
| target_dtype = first_leaf.dtype | ||
|
|
||
| def cvt(tensor, transpose=False): |
There was a problem hiding this comment.
can you rename this function to something more clear.
| ) | ||
| del pt_state_dict | ||
| gc.collect() | ||
| max_logging.log("Weight conversion complete!") |
There was a problem hiding this comment.
I think its best to validate params have been mapped correct. For example something like https://github.com/AI-Hypercomputer/maxdiffusion/blob/main/src/maxdiffusion/models/flux/util.py#L116
| max_logging.log(f"Loading VAE weights from: {safetensors_path}") | ||
| pt_state_dict = load_file(safetensors_path) | ||
|
|
||
| def get_w(key): |
There was a problem hiding this comment.
can you rename the function to something more clear.
| )(timesteps_proj.astype(dtype)) | ||
|
|
||
| if self.guidance_embeds and guidance is not None: | ||
| guidance_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(guidance) |
There was a problem hiding this comment.
where is 256 from? any way it can be passed instead of being hardcoded.
| @@ -0,0 +1,491 @@ | |||
| """ | |||
There was a problem hiding this comment.
is this model loaded replicated? Can it be sharded?
| # --------------------------------------------------------------------- | ||
| # PHASE A: Encode Prompt (Qwen3) | ||
| # --------------------------------------------------------------------- | ||
| print(f"[PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...") |
There was a problem hiding this comment.
for print statements, please use max_logging instead.
entrpn
left a comment
entrpn
left a comment
There was a problem hiding this comment.
left some comments + please change all print statements to max_logging instead.
3 participants
This branch contains all of the code for onboarding the Flux2.Klein model into this repo. PR contains support for both the 4B and 9B model variants.
Images can be generated with:
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.ymlFor the 9B model, point to the corresponding .yml file
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.ymlRunning this implementation on a v7-1 with FSDP=2 on Batch-size 8 (1024 by 1024) gives the following latency:
Some call-outs on the implementation:
src/maxdiffusion/generate_flux2klein.pyis the main user-facing file for running either the 4B and 9B files. It invokessrc/maxdiffusion/pipelines/flux/flux2klein_pipeline.py.src/maxdiffusion/models/flux/transformers/transformer_flux_flax.pyis heavily modified to add support for the new variants of the Attention blocks used in the Flux2.Klein models over the older generation of models.src/maxdiffusion/models/qwen3_flax.py