README.md

---
license: other
license_name: modified-mit
license_link: https://huggingface.co/AiArtLab/sdxs-1b/blob/main/LICENSE
pipeline_tag: text-to-image
datasets:
- CaptionEmporium/midjourney-niji-1m-llavanext
---

# Simple Diffusion XS (alpha version)

*XS Size, Excess Quality (not for now)* 

Train status: 2xRTX5090 in progress! / we need more gold / [donate please..](https://huggingface.co/AiArtLab/sdxs-1b#donations)

<img src="media/boy.png" height="256"/>

At AiArtLab, we strive to create a free, compact and fast model that can be trained on consumer graphics cards.

- Custom UNet: 1.6b parameters
- Qwen3.5:     1.8b parameters
- Custom VAE:  32ch8x16x
- Speed:       Sampling: 100%|██████████| 40/40 [00:01<00:00, 29.98it/s]
- Resolution:  Default 1024x1408, trained from 768 to 1408 with step 64 (384/704/32)
- Limitations: trained on small datasets ~1-2kk, focused on art / illustrations / anime (with ~10% photo / use "photo" in negative prompt)
- Captions:    danbooru, natural (short/medium), trained with 250 max toks

[Online demo on HF space](https://huggingface.co/spaces/LoveScapeAI/sdxs)

[comfy (unofficial)](https://github.com/customWF2026/CustomWFNodes)

### Key points

 - Dec 24: Started research on Linear Transformers.
 - Feb 25: Started research on UNet-based diffusion models.
 - Aug 25: Started research on different VAEs.
 - Sep 25: Created a simple VAE and a [vae collection](https://huggingface.co/AiArtLab/collections).
 - Dec 25: Trained SDXS-1B (0.8B at this moment), featuring an SD1.5-like UNet, Long CLIP, 16-channel simple VAE, and flow matching target.
 - Jan 26: Implemented a dual text encoder and SDXL-like style.
 - Feb 26: Reverted to classic architecture; tested all SDXL innovations and went back to simple diffusion.
 - Mar 26: Created an 32ch 8x/16x asymmetric VAE and switched to Qwen3.5 2B as text encoder.
 - Apr 26: Alpha version, finetuned VAE and UNet

### Samples with seed 0 and good prompts
![promo](media/result_grid.jpg)

### Text 2 image

```python
#!pip install -U torch torchvision
#!pip install -U diffusers accelerate transformers

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

pipe_id = "AiArtLab/sdxs-1b"
pipe = DiffusionPipeline.from_pretrained(
    pipe_id,
    torch_dtype=dtype,
    trust_remote_code=True
).to(device)

prompt = "1girl, smiling, red eyes, blue hair, white shirt"
refined = pipe.refine_prompts(prompt)[0]

negative_prompt = "bad quality grainy image with low details, incomplete text, despite numerous technical flaws and distorted figures"


image = pipe(
    prompt=refined,
    negative_prompt = negative_prompt,
    width = 1024,
    height = 1408,
    guidance_scale = 5,
    seed = 0,
).images[0]

image.show()
```

### Diversiry
Prompt: *cat*
![cat](media/cat.jpg)

Random prompts
![collage](media/gandr-collage.jpg)

### VAE

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](https://huggingface.co/AiArtLab/simplevae) with flux.1 quality based on aura vae.

```
Wan2.2-T2V-A14B            | MSE=8.908e-04 PSNR=33.21 LPIPS=0.052
FLUX.1-schnell VAE         | MSE=6.020e-04 PSNR=34.59 LPIPS=0.038
FLUX.2-klein-9B            | MSE=3.638e-04 PSNR=36.60 LPIPS=0.025
sdxs-1b                    | MSE=3.530e-04 PSNR=36.92 LPIPS=0.027

=== Percents ===
| Model                      |       MSE |      PSNR |     LPIPS |      Edge |
|----------------------------|-----------|-----------|-----------|-----------|
| Wan2.2-T2V-A14B            |      100% |      100% |      100% |      100% |
| FLUX.1-schnell VAE         |    148.0% |    104.2% |    136.7% |    128.8% |
| FLUX.2-klein-9B            |    244.8% |    110.2% |    209.7% |    165.9% |
| sdxs-1b                    |    252.4% |    111.1% |    193.1% |    172.1% |
```
To the best of our knowledge, this VAE probably achieves SOTA: [eval.py](src/eval.py)

### Image upscale (vae)
```
upscaled = pipe.image_upscale("media/girl.jpg")
upscaled[0].show()
```
[HF Demo](https://huggingface.co/spaces/LoveScapeAI/sdxs-1b-upscaler)

[Example](https://www.diffchecker.com/image-compare/EhOKWXFV/)

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](https://github.com/black-forest-labs/flux2/blob/main/src/flux2/autoencoder.py) and weights


| | |
|:---:|:---:|
| <img src="media/123456789.jpg" height="512"/> | <img src="media/123456789.png" height="512"/> |
| original | 2x upscale |

### true-to-source upscaling
| | | |
|:---:|:---:|:---:|
| <img src="media/up0.png" style="width:512px; height:512px; object-fit:cover;"/> | <img src="media/up1.png" style="width:512px; height:512px; object-fit:cover;"/> | <img src="media/up2.png" style="width:512px; height:512px; object-fit:cover;"/> |
| original | sdxs vae | seedvr2 |

### sdxs / resrgan
<img src="media/123.png" height="512"/>

### sdxs / swinir2
<img src="media/0000.png" height="512"/>

### upscale code example

```python
import torch
import numpy as np

from diffusers import AsymmetricAutoencoderKL
from typing import List, Union
from PIL import Image

vae = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae").cuda().half()
device = "cuda"

@torch.no_grad()
def image_upscale(
    image: Union[str, Image.Image, List[Union[str, Image.Image]]], 
    batch_size: int = 1
) -> List[Image.Image]:
    """
    Upscales images using asymmetric VAE (x2).
    Uses smart batching: processes in parallel if sizes match, else falls back to sequential.
    """
    images = [image] if isinstance(image, (str, Image.Image)) else image

    # 1. Preprocess: Load, Handle Alpha, Pad to %8, Normalize
    batch_data = []
    for img in images:
        if isinstance(img, str): img = Image.open(img)
        if img.mode == "RGBA":
            img = Image.alpha_composite(Image.new("RGBA", img.size, (255, 255, 255)), img)
        img = img.convert("RGB")

        w, h = img.size
        pw, ph = (8 - w % 8) % 8, (8 - h % 8) % 8
        if pw or ph:
            padded = Image.new("RGB", (w + pw, h + ph), (255, 255, 255))
            padded.paste(img)
            img = padded

        t = torch.from_numpy(np.array(img).astype(np.float32) / 127.5 - 1.0).permute(2, 0, 1)
        batch_data.append((t.to(device, torch.float16), w, h))

    # 2. Determine Execution Strategy
    # If all shapes are identical, use batch_size. Else fallback to 1.
    unique_shapes = {t.shape for t, _, _ in batch_data}
    step = batch_size if len(unique_shapes) == 1 else 1

    output_images = []

    # 3. Process Batches
    for i in range(0, len(batch_data), step):
        chunk = batch_data[i : i + step]

        # Stack tensors [B, C, H, W]
        tensors = torch.stack([c[0] for c in chunk])

        # Encode -> Decode (using mean for deterministic upscale)
        latents = vae.encode(tensors).latent_dist.mean
        decoded = vae.decode(latents.to(vae.dtype))[0]

        # 4. Post-process: Denormalize and Crop
        decoded = (decoded.clamp(-1, 1) + 1) / 2
        for j, tensor in enumerate(decoded):
            w, h = chunk[j][1], chunk[j][2] # Original sizes

            # Crop to exact 2x
            arr = tensor.cpu().permute(1, 2, 0).float().numpy()
            arr = arr[:h * 2, :w * 2]

            output_images.append(Image.fromarray((arr * 255).astype("uint8")))

    return output_images
    
up = image_upscale("1.jpg")
up[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.
![refine](media/refined.webp)

### Encode image (experimental)
```
negative_prompt = "red eyes"
emb, mask = pipe.encode_image("media/girl.jpg")
image = pipeline(
        prompt_embeds = emb,
        prompt_attention_mask = mask,
        negative_prompt = negative_prompt,
        guidance_scale = 4,
        width = 1088,
        height = 1344,
        seed = 0,
        batch_size = 1,
    )[0]
```


### Unet

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](https://wandb.ai/recoilme) of different configurations and trained dozens of [models](https://huggingface.co/AiArtLab/sdxs). 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](https://www.comet.com/recoilme/unet/356142c52c314078914d0c0db409e1f3?experiment-tab=images&viewId=new).
## 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
```
base_learning_rate = 4e-5
min_learning_rate = 4e-6

# 0.5 - pretrain (base forms)
# 1 - base train (composition)
# 3 - finetuning (anatomy)
# 5 - small details (faces)
base_learning_rate = base_learning_rate / learning_rate_scale
min_learning_rate = min_learning_rate / learning_rate_scale
```
## Tools and Optimization
The model comes with two scripts:

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](https://cloud.vast.ai/?ref_id=152324))
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).

## Acknowledgments
- **[Stan](https://t.me/Stangle)** — Key investor. Thank you for believing in us when others called it madness.
- **[the last neural cell](https://t.me/neural_cell)** - Thank you for providing 8xH100 for 48 hours 
- **Captainsaturnus**
- **Love. Death. Transformers.**
- **TOPAPEC**

## Datasets
- **[CaptionEmporium](https://huggingface.co/CaptionEmporium)**

## Donations

 - Rubles: [For users from Russia](https://www.tbank.ru/cf/90ensBQqpJj)
 - DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83
 - BTC: 3JHv9Hb8kEW8zMAccdgCdZGfrHeMhH1rpN
 - USTD
   - Ethereum / Polygon / BNB SmartChain: 0xD4388B6698dFaE1460E72099D4F208aaCA4f6E6C
   - Tron: TD7ey4h9igPGdcrcBcnZaz56R5tNgRZNvV
   - Solana: MMYFJeYEtYHrSNFHChytJDHbEDniXrnAxPNLhJ1LbkB

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

 - telegram [recoilme](https://t.me/recoilme) *prefered way
 - mail at aiartlab.org (slow response)

mail at aiartlab.org (slow response)

## Citation
```bibtex
@misc{sdxs,
  title={Simple Diffusion XS},
  author={recoilme, muinez and AiArtLab Team},
  url={https://huggingface.co/AiArtLab/sdxs-1b},
  year={2026}
}
```