The VAE in Simple Diffusion utilizes an asymmetric VAE architecture featuring an 8x encoder and a 16x decoder. While a compression factor of 8 is maintained during training, the resolution is effectively doubled during inference through an additional upscaling block. This strategy reduces training costs by an order of magnitude and boosts inference speed without perceptual quality loss. Effectively, this acts as an integrated latent upscaler. To ensure a fair comparison with other VAEs, we downsampled the generated images to match the input resolution for metric evaluation. The SDXS VAE was not trained from scratch but was initialized from weights of FLUX 2 VAE, then redisigned and retrained. We also trained 16 ch vae with flux.1 quality based on aura vae.
The Asymmetric VAE features a built-in 2x image upscaler. Because VAEs are trained to reconstruct images as accurately as possible, this upscaled acts as a "blind" processor: it enlarges the image by 2x within its trained range (512–768px) without altering its essence. Unlike model-based upscalers, it strictly preserves the original style and exact details without hallucinating new ones. This makes it ideal for precise, true-to-source upscaling. However, you will be disappointed if you expect AI magic—it cannot fix messy generations, invent missing textures, or turn a bad image into a masterpiece. This upscaler has been tested on video and can be used independently by downloading just the autoencoder model.
Development Note: We have not provided direct comparisons with other popular upscalers (like ESRGAN or SUPIR) because this is not the final version. The current model was trained in just 2 days on a single GPU. However, reaching this architecture required approximately 2 months of intensive research and significant personal funding. Currently, I am unemployed and unable to afford further training or compute to bring this to its final state, so I am releasing it as-is. I apologize for the lack of comparative benchmarks. The training code is fully open-source and available in our VAE collection: https://huggingface.co/collections/AiArtLab/vae. Based on Flux.2 autoencoder and weights
original
2x upscale
true-to-source upscaling
original
sdxs vae
seedvr2
sdxs / resrgan
sdxs / swinir2
upscale code example
python
1import torch
2import numpy as np
34from diffusers import AsymmetricAutoencoderKL
5from typing import List, Union
6from PIL import Image
78vae = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae").cuda().half()9device ="cuda"1011@torch.no_grad()12defimage_upscale(13 image: Union[str, Image.Image, List[Union[str, Image.Image]]],14 batch_size:int=115)-> List[Image.Image]:16"""
17 Upscales images using asymmetric VAE (x2).
18 Uses smart batching: processes in parallel if sizes match, else falls back to sequential.
19 """20 images =[image]ifisinstance(image,(str, Image.Image))else image
2122# 1. Preprocess: Load, Handle Alpha, Pad to %8, Normalize23 batch_data =[]24for img in images:25ifisinstance(img,str): img = Image.open(img)26if img.mode =="RGBA":27 img = Image.alpha_composite(Image.new("RGBA", img.size,(255,255,255)), img)28 img = img.convert("RGB")2930 w, h = img.size
31 pw, ph =(8- w %8)%8,(8- h %8)%832if pw or ph:33 padded = Image.new("RGB",(w + pw, h + ph),(255,255,255))34 padded.paste(img)35 img = padded
3637 t = torch.from_numpy(np.array(img).astype(np.float32)/127.5-1.0).permute(2,0,1)38 batch_data.append((t.to(device, torch.float16), w, h))3940# 2. Determine Execution Strategy41# If all shapes are identical, use batch_size. Else fallback to 1.42 unique_shapes ={t.shape for t, _, _ in batch_data}43 step = batch_size iflen(unique_shapes)==1else14445 output_images =[]4647# 3. Process Batches48for i inrange(0,len(batch_data), step):49 chunk = batch_data[i : i + step]5051# Stack tensors [B, C, H, W]52 tensors = torch.stack([c[0]for c in chunk])5354# Encode -> Decode (using mean for deterministic upscale)55 latents = vae.encode(tensors).latent_dist.mean
56 decoded = vae.decode(latents.to(vae.dtype))[0]5758# 4. Post-process: Denormalize and Crop59 decoded =(decoded.clamp(-1,1)+1)/260for j, tensor inenumerate(decoded):61 w, h = chunk[j][1], chunk[j][2]# Original sizes6263# Crop to exact 2x64 arr = tensor.cpu().permute(1,2,0).float().numpy()65 arr = arr[:h *2,:w *2]6667 output_images.append(Image.fromarray((arr *255).astype("uint8")))6869return output_images
7071up = image_upscale("1.jpg")72up[0].show()
Prompt refiner
refined = pipe.refine_prompts("1girl, solo, lazypos, cinematic lighting, volumetric lighting, dutch angle,
anime coloring, delicate, dreamy, natural shadow, serene atmosphere,
<lora:hololive_kaela_kovalskia_redebut:0.65> kaela20, blonde hair, red eyes, hair ribbon, red ribbon, long hair, half updo, tsurime,
BREAK
overalls, green overalls, white shirt, straw hat, hat flower, daisy, sleeves rolled up, three-quarter sleeves,
outdoors, flower field,
bouquet, holding bouquet, looking at viewer, smile,
m¨h0r_a97Q, refraction, prismatic light,
<lora:WSSKX_WAI:0.35>
<lora:ashima-il:0.35>
<lora:MYRHA-¨00:0.75>
<lora:Dramatic Lighting Slider:3.5>
<lora:spo_sdxl_10ep_4k-data_lora_webui:1.3>")
print(refined[0])
Refined prompt:
A blonde-haired Red Eyes girl with a hair ribbon, half-updo, and tsurime stands solo in a flower field holding a bouquet with a serene smile, wearing green overalls, a white shirt, rolled-up sleeves, and a straw hat with a flower while looking at the viewer under volumetric and natural lighting with a Dutch angle.
The UNet architecture in Simple Diffusion is a direct descendant and conceptual continuation of the ideas introduced in the first version of Stable Diffusion. By size it's 2xSD1.5. Key distinctions include a relatively small, yet sufficient, number of transformer blocks that ensure an even distribution of attention. Additionally, the number of channels in the final layer has been significantly increased to improve detail rendering. Overall, however, it remains a UNet, similar to SD 1.5.
Throughout the experiments, we tested hundreds of different configurations and trained dozens of models. Notably, we initially started from the SDXL architecture, assuming it would be a stronger baseline, but ultimately abandoned all the innovations proposed in it. These included uneven attention distribution with increased transformer block depth in the lower layers, a reduced number of blocks in the channel pyramid, micro-conditioning, the dual text encoder, text-time and so on. According to our experiments, all of these changes lead to increased training time and costs while having a near-zero or negative impact on the final result. In total, the investigation of various architectures and the search for the most efficient and optimal configuration took over a year.
Unfortunately, we were unable to secure grants for model training, with the exception of a grant from Google TPU—which, unfortunately, we were unable to utilize due to insufficient preparation and time constraints. As a result, training and experiments were financed primarily from our own funds and user donations. This left a significant mark on the model’s architecture.
We aimed to make it as small and cost-effective to train as possible while maintaining our quality generation requirements. So perhaps the limited budget even worked to our advantage.
Nevertheless, we remain hopeful for continued community support, which would allow us to further develop the model while remaining as independent as possible.
Text encoder
We tested various text encoders, including—but not limited to — CLIP, LongCLIP, SigLIP, MexmaSigLIP, Qwen3-0.6B, Qwen3-0.6B embeddings, Qwen3-1.7B. Ultimately, we settled on Qwen3.5-2B, which demonstrated unprecedented improvements in both quality and training speed.
We’d also like to highlight the LongCLIP model—its training speed is comparable to Qwen3.5, which is remarkable for its size.
Embeddings are extracted from the -2 layer, with a pooling layer from the last token as the first element. This was done to improve both composition and versatility — for example, it allows using a pooling layer from images instead of a textual instruction.
Training was conducted with a maximum of 250 tokens, and a 10% dropout rate was applied during training.
Additionally, the use of a full-fledged language model allowed us to integrate an optional prompt enhancement mechanism into the pipeline.
Retrospective and Key Takeaways
Research started in December 2024 after the release of the SANA model. We received a donation from Stan for fine-tuning SANA and, together with Stas, began fine-tuning and further developing it. Despite spending the entire budget, we did not achieve significant improvements. However, we were shocked by how poorly the model was trained and designed, and we became convinced that we could do better—though we were wrong.
Shifting Gears
By February 2025, we split our efforts and began designing our own architectures—which we are still doing today. Stas favored the DiT architecture, while I believed in UNet. Despite some differences in architectural views, we maintained close communication, shared our work, and supported each other throughout the process. We also engaged with the AIArtLab community (a virtual Telegram chat for those contributing to model development)—thank you all for your support.
Main mistake
One of my key mistakes was relying too heavily on LLMs and research papers. Research often presents minor improvements as groundbreaking innovations, and LLMs, trained on such content, can draw incorrect conclusions. From autumn 2025, I radically changed my strategy, switching to training simpler models (VAEs), where simple fine-tuning yielded more substantial improvements than expensive research projects—including fine-tuning a VAE to a quality level comparable to Flux-1 at the time.
This shift led me to adopt a zero-trust policy toward any external information not personally verified. This does not mean that you should not read papers, but I urge you not to trust the conclusions presented in them. This is an extremely radical approach, and I have intentionally radicalized it, but it allowed me to transition from reading papers and implementing other people's ideas to generating my own and training models.
As a result, I focused on building a strong local benchmark for rapid, cost-effective experiments on single rtx4080. This led me to train models on the "Butterflies" dataset—a set of 1,000 images of butterflies—where a model could be trained from scratch in just an hour to assess the impact of a hypothesis or improvement, example.
The Evolutionary Path
The second turning point was the transition to a continuous evolutionary improvement strategy. Unfortunately, the Butterflies dataset does not allow for evaluating prompt-following or anatomical generation capabilities. As a result, the model evolved incrementally rather than through revolutionary changes. The same model, from December 2025, underwent around 10 changes, including radical architectural shifts—while always preserving the pre-trained weights. It’s remarkable how well and quickly pre-trained models adapt to changes in architecture and external factors, even radical ones (e.g., switching VAE models, text encoders, or their combinations).
In addition to saving on training costs, this approach helped maintain minimal model size—for example, adding extra transformer blocks followed by an assessment of necessity and rolling back if the changes had no significant impact.
tldr;
Stop reading, start training
The Role of Hyperparameters
One of the initial mistakes was an excessive focus on hyperparameters during training. Ironically, 80% of training speed and quality depend on the model architecture (UNet) and the quality of embeddings (VAE), while other 20% is influenced by the text encoder’s embeddings. The rest is Role of Hyperparameters. The irony here is that Adam (adamw8bit) is surprisingly forgiving of hyperparameter errors, so I won’t even list them. Default is ok. I manipulate only LR at different train levels
A dataset script to convert a folder of image-text pairs into latent representations.
A training script provided as a single monolithic file.
Additionally, there’s a script that can be pasted directly into the terminal to automatically train the model with optimized parameters.
Training Optimization
All training was done using the AdamW8bit optimizer, which significantly reduced training costs.
I use 8xRTX5090 ($100/day on vast.ai)
For training from zero on my dataset rent with ~555Gb disk on ~10 days, copy create unet notebook and run with train lvl 0.5,1,2,3,4,5
This script will train with batch 64, with fp32(tf32) and must be stable.
Train:
apt update
apt install git-lfs
git config --global credential.helper store
git clone https://huggingface.co/AiArtLab/sdxs-1b
cd sdxs-1b
pip install -r requirements.txt -U
mkdir datasets
cd datasets
#hf download babkasotona/ds1234_noanime_704_vae8x16x --repo-type dataset --local-dir ds1234_noanime_704_vae8x16x
cd ..
nohup accelerate launch train.py &
Model Limitations:
Limited concept coverage due to the small dataset (1kk).
Thanks for your support! Every $5 equals a full day of RTX 4090 rental and makes the waifu 0.01% better!
Euro / dollars: sorry, we can’t accept payments via Patreon or Ko-fi. Please register on vast.ai (this is the GPU provider we use for training the model), top up any amount, and transfer: Transfer Money / User / Email: vadim-kulibaba@yandex.ru.
