madebyollin/taeqi2_1

🍰 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.

This repo contains .safetensors versions of the TAEQI2.1 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)