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

zeromodels/qwen-image

Paper: Qwen-Image Technical Report (arXiv:2508.02324)

Pure-Keras 3 conversion of Qwen/Qwen-Image for zeromodels. One implementation runs unmodified on TensorFlow / Torch / JAX. The container ships the 60-layer double-stream MMDiT, the Wan-derived 16-channel VAE (8× spatial), and the Qwen2.5-VL text tower in model.weights.json shards (~53 GiB at bfloat16), plus zm_config.json (component configs, the checkpoint's FlowMatchEulerDiscreteScheduler with dynamic shifting, and default generation settings) and tokenizer.json (Qwen2 BPE under the Diffusers ChatML template). Weights are stored in bfloat16 and load in bfloat16 by default; pass load_dtype="float32" to from_weights for a float32 model. This checkpoint targets 1024×1024 images (packed 2×2 latents over a 128×128×16 grid).

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

Architecture

Component	zeromodels class	Details
Denoiser	QwenImageTransformer2DModel	60 double-stream MMDiT layers; packed (B, H/2·W/2, 64) tokens
Autoencoder	AutoencoderKLQwenImage	Wan-style KL VAE, z_dim 16, 8× spatial, latents_mean / latents_std
Text encoder	QwenImageTextEncoderModel	Qwen2.5-VL text tower (28 layers, 3584-d)
Scheduler	FlowMatchEulerDiscreteScheduler	rectified flow, dynamic resolution shifting from packed sequence length
Tokenizer	QwenImageTokenizer	Qwen2 BPE + Diffusers ChatML system prompt

Quick start

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

from PIL import Image
from zeromodels.models.qwen_image import QwenImageTextToImage, QwenImageTokenizer

model = QwenImageTextToImage.from_weights("zeromodels/qwen-image")
tokenizer = QwenImageTokenizer.from_weights("zeromodels/qwen-image")

inputs = tokenizer("a photo of a capybara wearing a wizard hat")
images = model.generate(**inputs, num_inference_steps=50, guidance_scale=4.0, seed=0)
Image.fromarray(images[0]).save("capybara.png")  # (1024, 1024, 3) uint8

generate takes the tokenizer's input_ids / attention_mask, optional negative_input_ids when guidance_scale > 1, num_inference_steps, guidance_scale, a seed, or explicit packed latents; height / width default to 1024.

Variant	Hub	Notes
qwen-image	zeromodels/qwen-image	1024px, 50 flow-match steps, guidance_scale=4.0

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 16px.
• QwenImageModel.from_weights(...) loads the bare container without the generation loop.
• On-the-fly hf: conversion is not supported for diffusion models.
• See the Qwen-Image docs.

License

The weights are redistributed under the Apache License 2.0 of the upstream checkpoint. By using them you agree to those terms.

Notice

Modifications by zeromodels (https://github.com/IMvision12/ZeroModels): the checkpoint released at https://huggingface.co/Qwen/Qwen-Image 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, and to the Hugging Face diffusers team, whose implementation this port was verified against.