🍰 Tiny AutoEncoder for Qwen Image 2.1

TAEQI2.1 is very tiny autoencoder which uses the same "latent API" as Qwen Image 2.1's VAE. TAEQI2.1 is useful for real-time previewing of the Qwen Image 2.1 generation process, as well as general resource-constrained encoding/decoding. Like the Qwen Image 2.1 VAE, TAEQI2.1 uses 16x spatial compression, 64 latent channels, and RGBA images. TAEQ2.1-Tinier is an even-lower-parameter-count variant of TAEQI2.1 for parameter-constrained usecases (way fewer parameters, slightly faster speed, slightly reduced quality).

This repo contains .safetensors versions of the TAEQI2.1 and TAEQI2.1-Tinier weights.

TAEQI2.I was Built with Qwen (in the sense that the TAEQI2.1 encoder was trained to mimic outputs from the Qwen-Image-2.1 VAE encoder).

Using in 🧨 diffusers

NOTE: Like TAEF2, TAEQI2.1's architecture isn't properly integrated into Diffusers yet. So for now you'll want some wrapper code:

pip install git+https://www.github.com/huggingface/diffusers # needed for Qwen Image 2.1 support
wget -nc -nv https://raw.githubusercontent.com/madebyollin/taesd/refs/heads/main/taesd.py -O taesd.py
wget -nc -nv https://huggingface.co/madebyollin/taeqi2_1/resolve/main/taeqi2_1.safetensors -O taeqi2_1.safetensors

# Construction
from taesd import TAESD
import torch
import safetensors.torch as stt
from diffusers.utils.accelerate_utils import apply_forward_hook

def convert_diffusers_sd_to_taesd(sd):
    out = {}
    for k, v in sd.items():
        encdec, _layers, index, *suffix = k.split(".")
        offset = 0
        if encdec == "decoder":
            offset = +1
        out[".".join([encdec, str(int(index)+offset), *suffix])] = v
    return out

class DotDict(dict):
    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__

class DiffusersTAEQI21Wrapper(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dtype = torch.bfloat16
        self.taesd = TAESD(encoder_path=None, decoder_path=None, latent_channels=64, arch_variant="f16", image_channels=4).to(self.dtype)
        self.taesd.load_state_dict(convert_diffusers_sd_to_taesd(stt.load_file("taeqi2_1.safetensors")))
        # TAEQI2.1 consumes / produces normalized latents directly, so the pipeline's latent scale / shift should be a no-op
        self.config = DotDict(z_dim=64, latents_mean=[0.0] * 64, latents_std=[1.0] * 64)

    @apply_forward_hook
    def encode(self, x):
        # x is (B, 4, 1, H, W) RGBA in [-1, 1]; latents are (B, 64, 1, H/16, W/16)
        x = x.squeeze(2)
        latents = self.taesd.encoder(x.to(self.dtype).mul(0.5).add_(0.5)).to(x.dtype).unsqueeze(2)
        return DotDict(latent_dist=DotDict(sample=lambda generator=None: latents, mode=lambda: latents))

    @apply_forward_hook
    def decode(self, x, return_dict=True):
        # x is (B, 64, 1, H/16, W/16); output is (B, 4, 1, H, W) RGBA in [-1, 1]
        x = self.taesd.decoder(x.squeeze(2).to(self.dtype)).mul(2).sub_(1).clamp_(-1, 1).to(x.dtype).unsqueeze(2)
        return dict(sample=x) if return_dict else (x,)

taeqi2_1_diffusers = DiffusersTAEQI21Wrapper().eval().requires_grad_(False)

# Usage
from diffusers import QwenImage21Pipeline

device = "cuda"
dtype = torch.bfloat16

pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=dtype)
pipe.vae = taeqi2_1_diffusers
pipe.enable_model_cpu_offload() # pipe = pipe.to(device)

prompt = "A slice of delicious New York-style berry cheesecake"
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    generator=torch.Generator(device="cpu").manual_seed(0)
).images[0]
image.save("qwen-image-2.1.png")
image

Quality comparisons (see thread)