[new-model] Port Z-Image T2I to FastVideo

[mrinaald]

Summary

This PR introduces the initial Z-Image text-to-image integration work into FastVideo, focused on core component compatibility and parity validation against the local Z-Image reference implementation. It adds native Qwen3 text-encoder support, aligns scheduler timestep behavior with Z-Image reference semantics, and adds local parity tests for scheduler, tokenizer, text encoder, and VAE decode path. This is an in-progress port intended to establish correctness foundations before final end-to-end pipeline parity.

What changed

Model (Text Encoder)

• Added native Qwen3 encoder config and model implementation in FastVideo.
• Implemented Qwen3 architecture wiring with RoPE, grouped KV attention handling, SDPA attention path, RMSNorm, and tensor-parallel compatible linear layers.
• Added robust weight loading for Qwen3 checkpoints, including support for model.-prefixed keys and stacked parameter remapping (QKV and gate/up projections).
• Exported Qwen3 config in encoder config registry and added model registry mappings for Qwen3Model and Qwen3ForCausalLM.

Scheduler

• Extended FlowMatchEulerDiscreteScheduler with use_reference_discrete_timesteps.
• Added Z-Image-compatible timestep construction mode (build num_steps + 1 linspace and drop terminal point) to match reference scheduler behavior.
• Preserved existing scheduler behavior as default when the new flag is not enabled.

Parity tests (local)

• Added scheduler parity test against Z-Image reference scheduler:
  • default timestep schedule + step loop parity
  • dynamic shifting parity with mu
• Added tokenizer parity test:
  • TokenizerLoader vs AutoTokenizer parity on input_ids and attention_mask
  • chat template parity checks where available
• Added VAE decode parity test:
  • FastVideo AutoencoderKL vs Z-Image reference autoencoder using identical local weights/config.
• Added Qwen3 text encoder parity test:
  • FastVideo Qwen3 vs transformers AutoModel on identical local checkpoint/tokenization
  • compares last hidden state and hidden_states[-2] on valid-token positions

How to test

Prerequisites

• Local Z-Image reference repo available at Z-Image/src
• Local model assets under official_weights/Z-Image
• PyTorch + transformers + safetensors installed

From repo root, run:

pytest tests/local_tests/zimage/

Current status / known gap

• Scheduler parity: passing
• Tokenizer parity: passing
• VAE decode parity: passing
• Qwen3 encoder parity: passing in float32 path, not passing in bfloat16
• ZImageTransformer2DModel: in-progress

[github-actions]

Welcome to FastVideo! Thanks for your first pull request.

How our CI works:

PRs run a two-tier CI system:

1. Pre-commit — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
2. Fastcheck — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
3. Full Suite — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the ready label.

Before your PR is reviewed:

• pre-commit run --all-files passes locally
• You've added or updated tests for your changes
• The PR description explains what and why

If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.

Useful links:

• Contributing Guide
• Development Roadmap
• Slack

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success~=pre-commit

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success~=pre-commit
• check-success=fastcheck-passed
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces support for the Qwen3 model, including its configuration, architecture implementation, and registration within the model registry. It also updates the FlowMatchEulerDiscreteScheduler to include an option for reference discrete timesteps, ensuring parity with specific implementations like Z-Image. Additionally, a suite of parity tests for the encoder, scheduler, tokenizer, and VAE has been added. Feedback focuses on ensuring architectural correctness for causal masking when the attention mask is missing and improving the robustness of the weight loading logic by using safer attribute access.

[gemini-code-assist]

The is_causal parameter is hardcoded to False. Since Qwen3 is a causal architecture (as indicated by its registration as Qwen3ForCausalLM), it should use causal masking. While sdpa_mask handles this when an attention_mask is provided, if attention_mask is None, the model will incorrectly default to bidirectional attention. Consider setting is_causal=True when attention_mask is None to maintain architectural correctness.

[gemini-code-assist]

Accessing param.weight_loader directly may raise an AttributeError if the parameter belongs to a standard layer that hasn't been extended with a custom loader. It is safer to use getattr(param, "weight_loader", default_weight_loader), consistent with the implementation in the else block at line 426.

Suggested change

	weight_loader(param, loaded_weight, shard_id)
	weight_loader = getattr(param, "weight_loader", default_weight_loader)

[mergify]

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