Run Qwen-Image-2.1 with Keras 3: JAX, PyTorch, or TensorFlow

zeromodels/qwen-image-2.1

Pure-Keras 3 conversion of Qwen/Qwen-Image-2.1 for zeromodels. One implementation runs unmodified on TensorFlow / Torch / JAX. The container ships the 32-layer single-stream DiT, the residual 64-channel RGBA VAE (16× spatial), and the Qwen3-VL text tower in model.weights.json shards (~28 GiB at bfloat16), plus zm_config.json (component configs, the checkpoint's FlowMatchEulerDiscreteScheduler with dynamic shifting, and default generation settings) and tokenizer.json (Qwen3 BPE under the Diffusers ChatML template). Weights are stored in bfloat16, the checkpoint's native precision, and load in bfloat16 by default; pass load_dtype="float32" to from_weights for a float32 model. This checkpoint targets 1024×1024 images (a 64×64×64 unpatched latent).

For model details, intended use and limitations, see the upstream model card.

Architecture

Component	zeromodels class	Details
Denoiser	QwenImage21Transformer2DModel	32 single-stream block-causal DiT layers; unpatched (B, H·W, 64) tokens; mixed-granularity attention; causal_condition text modulation
Autoencoder	AutoencoderKLQwenImage21	residual RGBA VAE, z_dim 64, 16× spatial compression, latents_mean / latents_std normalization
Text encoder	QwenImage21TextEncoderModel	Qwen3-VL text tower (36 layers, 4096-d, 32 heads / 8 KV); pre-norm features (before final RMSNorm)
Scheduler	FlowMatchEulerDiscreteScheduler	rectified flow, dynamic resolution shifting (mu from image sequence length), shift_terminal
Tokenizer	QwenImage21Tokenizer	Qwen3 BPE + ChatML system prompt (Comprehend and analyze the provided prompt.)

Quick start

import os
os.environ["KERAS_BACKEND"] = "torch"  # or "jax" / "tensorflow"

from PIL import Image
from zeromodels.models.qwen_image_21 import QwenImage21TextToImage, QwenImage21Tokenizer

model = QwenImage21TextToImage.from_weights("zeromodels/qwen-image-2.1")
tokenizer = QwenImage21Tokenizer.from_weights("zeromodels/qwen-image-2.1")

inputs = tokenizer("a neon shop sign that reads \"QWEN IMAGE 2.1\", rainy night")
images = model.generate(**inputs, num_inference_steps=40, guidance_scale=1.0, seed=0)
Image.fromarray(images[0]).save("sign.png")  # (1024, 1024, 3) uint8

generate takes the tokenizer's input_ids / attention_mask (batch them for several prompts), an optional tokenized negative prompt (negative_input_ids) when guidance_scale > 1, num_inference_steps, guidance_scale, a seed, or explicit unpatched latents of shape (batch, H·W, 64); height / width default to 1024.

Variant	Hub	Notes
qwen-image-2.1	zeromodels/qwen-image-2.1	1024px, 40 flow-match steps, guidance_scale=1.0 (true CFG optional)

Tips

• Set KERAS_BACKEND before importing Keras / zeromodels.
• The graphs are built for 1024px. Pass transformer_sample_size=, vae_sample_size= to from_weights to build for another multiple of 32px (weights are resolution-independent).
• QwenImage21Model.from_weights(...) loads the same repo as the bare container (.transformer / .vae / .text_encoder) without the generation loop.
• Both channels_last and channels_first are supported (keras.config.set_image_data_format before loading); generate always returns (batch, H, W, 3) uint8 (VAE RGBA cropped to RGB).
• On-the-fly hf: conversion is not supported for diffusion models; the checkpoints are hosted here, converted once.
• See the Qwen-Image-2.1 docs.

License

The weights are redistributed under the Qwen Research License of the upstream checkpoint (non-commercial research / evaluation use unless you obtain a separate commercial license from Qwen). By using them you agree to those terms. Include a copy of LICENSE with any redistribution.

Notice

Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 Hangzhou Tongyi Laboratory Technology Co., Ltd. All Rights Reserved.

Modifications by zeromodels (https://github.com/IMvision12/ZeroModels): the checkpoint released at https://huggingface.co/Qwen/Qwen-Image-2.1 was converted to the Keras 3 weights layout of zeromodels (model.weights.json, model_*.weights.h5, zm_config.json, tokenizer.json), stored in bfloat16. The model architecture and the parameter values are unchanged; the weight names and the file format differ from the release.

Special Thanks

Thank you to the Qwen team for training and releasing Qwen-Image-2.1, and to the Hugging Face diffusers team, whose implementation this port was verified against.