Zen Image Edit

Qwen-Image-2.1 on a 0.8B text encoder. Text-to-image, character and scene editing, and transparent (RGBA) generation in one pipeline.

What changed

The bundled VAE is not Qwen-Image-2.1's. It is a finetuned decoder, from madebyollin/texture-fix-vae-for-qwen-image-2.1, trained to remove a 2 px lattice the stock decoder leaves on its output. The encoder is bit-identical, so the latent space is unchanged. The original is at Qwen/Qwen-Image-2.1/vae. The file is fp16 and the pipeline loads it in fp16: rounding the weights to fp16 costs nothing measurable (PSNR −0.0002 dB, zero pixel difference). Measurements below.

Low-VRAM: recoilme/zen-image-edit-gguf — the transformer in GGUF Q5_K, 13.55 → 4.52 GiB, so it stays resident on a 16 GB card (8.9 GB peak) instead of being offloaded every step. text_fusion is left unquantized there.

The text encoder is replaced by Qwen3.5-0.8B plus a 158M adapter, fine-tuned to reproduce what the native encoder produced — both from plain text and from text read together with the reference images (Improved using Qwen). The adapter lives inside the DiT as its text-fusion block, so the whole model is one self-contained diffusers folder and no 17.5 GB encoder is needed anywhere. The sampler runs a plain static shift of 5.0 instead of the original dynamic shifting.

The bundled adapter is revision v12: its attention-branch position table covers 2304 slots and it was fine-tuned at the real inference geometry (~2000-token conditions at 1024 px), so reference images keep their positions instead of falling into a zero-padded tail — the vision cosine against the native encoder moved 0.93 → 0.97, text stayed at 0.95.

The finetuned VAE

The stock decoder upsamples four times with nearest-exact x2 + Conv2d 3x3, and a 3-tap kernel cannot undo the phase nearest-exact locks in -- so a regular 2 px lattice ends up fixed to the output stride, identical on every image.

Trained decoder-only, 5300 steps, on top of the usual lpips / mse / mae / edge mix plus one extra term: the peak-to-peak difference between the reconstruction's and the target's 2x2 sublattice means. Content reproduced faithfully contributes zero, so only the added lattice is penalised. The 45 decoder tensors texture-fix trained were frozen.

Reconstruction, 32 held-out images, 512:

Generation, 768x1280, 50 steps, fixed seed, only the VAE swapped -- three prompts with faces:

The change is invisible to the eye at 100%: a 0.5-2% contrast pattern at ~2 cycles/mm. Hence measured, not eyeballed.

Examples

Every image below is generated by this pipeline with 30 steps at 1024 px.

Text-to-image

Edit — one condition image (background change, subject kept)

Edit — two condition images (character replacement: <image1> is the edit target and keeps its pose, clothing and scene; the identity is copied from <image2>)

Edit — three condition images (target and composition from <image1>, the person from <image2>, colour and lighting from <image3>)

Transparent RGBA

Usage

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("AiArtLab/zen-image-edit", custom_pipeline="pipeline",
                                         trust_remote_code=True, dtype=torch.float16)
pipe.enable_model_cpu_offload()   # 14.5 GB DiT + fp32 VAE decoder do not co-reside on 32 GB

# text-to-image
image = pipe(prompt="a red fox in a snowy forest at dusk, cinematic, 85mm",
             output_resolution=1024, num_inference_steps=30,
             generator=torch.Generator("cuda").manual_seed(1234)).images[0]

# editing: 1..N condition images. The FIRST one is the edit target, the rest are references;
# reference them in the prompt by TAG <image1>, <image2>, ...
image = pipe(prompt="Replace the woman in <image1> with the woman from <image2>; keep <image1> pose, "
                    "clothing and background unchanged.",
             image=[scene_image, ref_image],
             output_resolution=1024, num_inference_steps=30,
             generator=torch.Generator("cuda").manual_seed(1234)).images[0]

Editing convention: the first image is the one being edited (<image1>), everything after it is a reference. That is the model's own convention and what the stock ComfyUI node documents; feeding the reference first is the usual reason a swap "does not happen" (the model then edits the reference). Note that the canvas size still comes from the last image's aspect ratio — pass height/width explicitly to pin it.

custom_pipeline="pipeline" builds the shipped pipeline.py and trust_remote_code=True lets it run, so no clone is needed. (_class_name is kept a plain string in model_index.json because that is what Hub tooling expects; the [file, class] form diffusers also accepts makes the Hub print a configuration warning.) Cloning works too and gives the class directly:

from pipeline import ZenImageEditPipeline
pipe = ZenImageEditPipeline.from_pretrained(".", dtype=torch.float16)

CLI — one image, or a whole file of prompts (one per line, # starts a comment, blank lines are skipped; the pipeline is loaded once for the whole file):

python example.py --prompt "a red fox in a snowy forest" --out fox.png
python example.py --prompts-file prompts.txt --out gens --size 1024 --steps 30
python example.py --prompt "..." --width 1280 --height 768 --out wide.png
python example.py --prompt "..." --negative "low quality, blurry, watermark" --cfg 3 --out cfg.png
python example.py --prompt "..." --scheduler-test --shift 5 --out ab.png

--scheduler-test renders every prompt twice with the same seed — the shipped static --shift (5.0) and Qwen-Image-2.1's original dynamic-shift schedule — and glues the pair with labels, so a schedule change can be judged without rerunning anything by hand.

--size sets a square frame (or the frame area when --image supplies the aspect ratio); --width/--height override it and are floored to a multiple of 32. --cfg is true_cfg_scale and defaults to 1.5 with a default negative prompt. Stock Qwen-Image-2.1 runs without guidance (--cfg 1), which is still what --cfg 1 and negative_prompt=None give you.

Requirements: torch, transformers, accelerate and a diffusers built with Qwen-Image-2.1 (pip install git+https://github.com/huggingface/diffusers) — the transformer subclasses QwenImage21Transformer2DModel. trust_remote_code saves the clone, it does not save the 17 GB of weights.

ComfyUI

The same adapter runs in ComfyUI, also without the 17.5 GB encoder — nodes, a ready-made workflow and the adapter file are in recoilme/zen-image-edit-comfyui.

• workflow: workflows/zen-image-edit_ui.json
• adapter for the loader node: releases/v1

Files

pipeline.py          ZenImageEditPipeline — one class for t2i and editing, as QwenImage21Pipeline
transformer.py       QwenImage21FusionTransformer2DModel + the text-fusion blocks
example.py           CLI for both modes
transformer/         DiT config + 2 fp16 shards, adapter merged in as text_fusion.*
text_encoder/        Qwen3.5-0.8B, fp16
processor/           its processor (image slicing + tokenization)
tokenizer/           its tokenizer
vae/                 finetuned Qwen-Image-2.1 VAE decoder, fp16 (see above)
scheduler/           FlowMatchEulerDiscreteScheduler config
media/               the examples above

QwenImage21FusionTransformer2DModel is a custom class defined in transformer.py, not registered inside diffusers, so the pipeline publishes it on the diffusers module at import time. That is what makes the trust_remote_code=True one-liner above work; without it the stock component loader would not find the DiT class.

Limitations

• English only — that is all the adapter was trained and tested on; other languages drift.
• Numerals on signage come out wrong: "OPEN 24 HOURS" renders as "OPEN 26 HOURS" on every seed tried. Words are fine.
• Non-photo references transfer less faithfully than photographic ones: the adapter imitates the native encoder, so its ceiling is the native encoder's ceiling.
• Batch size >1 at 1024 px peaks near 28 GB; one prompt per call is the safe mode.

NOTICE

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

This is a derivative work of Qwen-Image-2.1 — the full agreement is in LICENSE, the list of modified files and the remainder of the required attribution is in NOTICE. The Qwen3.5-0.8B text encoder is redistributed under the Apache License 2.0, see LICENSE-Qwen3.5-0.8B.

Contacts

Please contact with us if you may provide some GPU's or money on training

• telegram recoilme *prefered way
• mail at aiartlab.org (slow response)

Citation

@misc{zenimageedit,
  title={Zen Image Edit},
  author={recoilme and AiArtLab Team},
  url={https://huggingface.co/AiArtLab/zen-image-edit},
  year={2026}
}