[MAX] Add UMT5 text encoder for Wan diffusion

[jglee-sqbits]

Stacked PRs:

• [MAX] Add Wan video generation examples and comparison benchmark #19
• [MAX] Add Wan I2V diffusion pipeline #18
• [MAX] Add Wan T2V diffusion pipeline with MoE support #17
• [MAX] Add Wan transformer model with block-level compilation #16
• [MAX] Add Wan VAE and refactor autoencoder module #15
• ->[MAX] Add UMT5 text encoder for Wan diffusion #14
• [MAX] Add UniPC multistep scheduler for Wan diffusion #13

---

[MAX] Add UMT5 text encoder for Wan diffusion

Summary

Add a MAX-native UMT5 text encoder for Wan diffusion pipelines.

Description

• Implements the UMT5 encoder architecture using max.nn (Module V2 graph API)
• Supports float32 → bfloat16 weight casting via WeightData.astype() for Wan checkpoints that store text encoder weights in float32
• Includes T5-style relative position bias and gated GeLU feed-forward
• Handles diffusers weight key remapping (e.g. shared.weight alias dedup)

UMT5 is the text encoder used by all Wan 2.1/2.2 models. It produces 4096-dim text embeddings consumed by the Wan transformer.

Dependencies

None — can be merged independently.

Checklist

• PR is small and focused
• I ran ./bazelw run format to format my changes

Assisted-by: Claude Code

Assisted-by: Claude Code

[gemini-code-assist]

Code Review

This pull request introduces the UMT5 encoder architecture, providing the necessary components for model configuration, core neural network layers, and weight adaptation. The implementation includes a UMT5Model component designed for the MAX pipeline. Feedback suggests several improvements for better flexibility and performance: avoiding hardcoded data types to support various precisions, pre-calculating constant values like the attention scale and data type minimums outside the forward pass to optimize graph construction, and moving configuration parsing logic into the configuration class to prevent side effects from in-place mutation.

[gemini-code-assist]

Hardcoding DType.bfloat16 here overrides the encoding and config settings provided during initialization. While the comment explains this is a workaround for specific Wan 2.1 checkpoints, it limits the reusability of the UMT5Model component for other models that might require different precision (e.g., float32 or float16). This precision-forcing logic should ideally be handled during configuration resolution or guarded by a check to ensure it only applies when necessary.

[gemini-code-assist]

The scale value is constant for a given model configuration and does not depend on the input sequence length. It should be pre-calculated in the __init__ method of UMT5Attention to avoid redundant math.log calls and arithmetic operations during graph construction for every layer.

[gemini-code-assist]

The _DTYPE_MIN dictionary is defined inside the __call__ method, which means it is re-allocated on every forward pass during graph construction. This is inefficient. Please move this dictionary to the module level or as a class attribute.

[gemini-code-assist]

The UMT5EncoderModel constructor is mutating the config object in-place to store parsed activation information. This is a side effect that can affect other parts of the application using the same configuration instance. It is recommended to move this parsing logic into the UMT5ConfigBase class (e.g., as properties or using a Pydantic validator) to keep the model implementation clean and side-effect free.