512ep20

README.md

@@ -7,6 +7,19 @@ pipeline_tag: text-to-image
 
 *XS Size, Excess Quality*
 
+At AiArtLab, we aim to develop a compact (1.7b) and fast (3sec/image) model that can be trained on consumer-grade graphics cards, all while operating on a limited budget. 
+
+
+We have chosen the multilingual encoder Mexma-SigLIP, which supports 80 languages and processes entire sentences rather than individual tokens. Our chosen VAE architecture, AuraDiffusion, preserves details and anatomy without the blurring effects seen in other models.
+
+
+For training, we use AdamW-8bit, which allows for larger batch sizes and accelerates training on cost-effective GPUs. Our model has been trained on approximately one million images with various resolutions and styles, including anime and realistic photos. We employed a variety of annotation methods, combining both manual and automated approaches.
+
+
+However, our model does have some limitations:
+- Limited concept coverage due to the small dataset size.
+- The Image2Image functionality requires further training.
+
 
 Train status, in progress: [wandb](https://wandb.ai/recoilme/unet)
 
@@ -107,7 +120,7 @@ if __name__ == "__main__":
     scheduler = DDPMScheduler.from_pretrained(pipeid, subfolder="scheduler")
 
 
-    height, width = 384, 384
+    height, width = 576, 384
     num_inference_steps = 40
     output_folder, project_name = "samples", "sdxs"
     latents = generate_latents(
@@ -126,28 +139,6 @@ if __name__ == "__main__":
     print("Images generated and saved to:", output_folder)
 ```
 
-## Introduction
-*Fast, Lightweight & Multilingual Diffusion for Everyone*
-
-We are **AiArtLab**, a small team of enthusiasts with a limited budget. Our goal is to create a compact and fast model that can be trained on consumer graphics cards (full training cycle, not LoRA). We chose U-Net for its ability to efficiently handle small datasets and train quickly even on a 16GB GPU (e.g., RTX 4080). Our budget was limited to a few thousand dollars, significantly less than competitors like SDXL (tens of millions), so we decided to create a small but efficient model, similar to SD1.5 but for 2025 year.
-
-## Encoder Architecture (Text and Images)
-We experimented with various encoders and concluded that large models like LLaMA or T5 XXL are unnecessary for high-quality generation. However, we needed an encoder that understands the context of the query, focusing on "prompt understanding" versus "prompt following." We chose the multilingual encoder Mexma-SigLIP, which supports 80 languages and processes sentences rather than individual tokens. Mexma accepts up to 512 tokens, creating a large matrix that slows down training. Therefore, we used a pooling layer to simplify 512x1152 matrix with plain 1x1152 vector. Specifically, we passed it through a linear model/text projector to achieve compatibility with SigLIP embeddings. This allowed us to synchronize text embeddings with images, potentially leading to a unified multimodal model. This functionality enables mixing image embeddings with textual descriptions in queries. Moreover, the model can be trained without text descriptions, using only images. This should simplify training on videos, where annotation is challenging, and achieve more consistent and seamless video generation by inputting embeddings of previous frames with decay. In the future, we aim to expand the model to 3D/video generation.
-
-## VAE Architecture
-We chose an unconventional 8x 16-channel AuraDiffusion VAE, which preserves details, text, and anatomy without the 'haze' characteristic of SD3/Flux.
-
-## Training Process
-### Optimizer
-We tested several optimizers (AdamW, Laion, Optimi-AdamW, Adafactor, and AdamW-8bit) and chose AdamW-8bit. Optimi-AdamW demonstrated the smoothest gradient decay curve, although AdamW-8bit behaves more chaotically. However, its smaller size allows for larger batch sizes, maximizing training speed on low-cost GPUs (we used 4xA6000 and 5xL40s for training).
-
-### Dataset
-We trained the model on approximately 1 million images: 60 epochs on ImageNet at 256 resolution (wasted time because of low-quality annotations) and 8 epochs on CaptionEmporium/midjourney-niji-1m-llavanext, plus realistic photos and anime/art at 576 resolution. We used human prompts, Caption Emporium provided prompts, WD-Tagger from SmilingWolf, and Moondream2 for annotation, varying prompt length and composition to ensure the model understands different prompting styles. The dataset is extremely small, leading the model to miss many entities and struggle with unseen concepts like 'a goose on a bicycle.' The dataset also included many waifu-style images, as we were interested in how well the model learns human anatomy rather than drawing 'The Astronaut on horseback' skills. While most descriptions were in English, our tests indicate the model is multilingual.
-
-## Limitations
-- Limited concept coverage due to the extremely small dataset.
-- The Image2Image functionality needs further training (we reduced the SigLIP portion to 5% to focus on text-to-image training).
-
 ## Acknowledgments
 - **[Stan](https://t.me/Stangle)** — Key investor. Primary financial support - thank you for believing in us when others called it madness.
 - **Captainsaturnus** — Material support.

src/cherrypick.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ad16f012308a461f220b8b9330c5606f5d1d25a8fac8cd64f79d29358079d98
-size 15687
+oid sha256:c3bc29d6c8ede5a64c8ae9b2f8f824d4edfe91209bc6a6363a43ab66ec01d68f
+size 44788

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:92d6a1b7034bd2de712560bf42723551e6884a9f3c379855329e50dee422c889
-size 1401708
+oid sha256:d35e0b0f890df0408db1efd1fe1daed6b8b27b3021e90fa10ded72d702f26677
+size 4885422

train.py

@@ -28,9 +28,9 @@ import torch.nn.functional as F
 ds_path = "datasets/576"
 project = "unet"
 batch_size = 50
-base_learning_rate = 5e-5
+base_learning_rate = 2e-5
 min_learning_rate = 1e-5
-num_epochs = 19
+num_epochs = 5
 # samples/save per epoch
 sample_interval_share = 5
 use_wandb = True
@@ -50,7 +50,7 @@ torch.backends.cuda.enable_mem_efficient_sdp(True)
 dtype = torch.float32
 save_barrier = 1.03
 dispersive_temperature=0.5
-dispersive_weight=0.1
+dispersive_weight=0.05
 percentile_clipping = 97 # Lion
 steps_offset = 1 # Scheduler
 limit = 0
@@ -628,7 +628,7 @@ def create_optimizer(name, params):
     if name == "adam8bit":
         import bitsandbytes as bnb
         return bnb.optim.AdamW8bit(
-            params, lr=base_learning_rate, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01
+            params, lr=base_learning_rate, betas=(0.9, 0.995), eps=1e-7, weight_decay=0.001
         )
     elif name == "adam":
         return torch.optim.AdamW(

unet/diffusion_pytorch_model.fp16.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ba82fba0a530fa8d0a310e5946b98c08bb5e482aa26c65a74fed55e59b27a99c
-size 7014306128
+oid sha256:5bb9c77df1219662fc14575e71304875d9a98861cffe930e83483e4274a34238
+size 3507231768

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:faa8982e546e7c861e550b58db204dfbf3a67688ed0b7b06b10fd0c58a596955
+oid sha256:d16cb10d5304bc2214ed9559993089a98cd6696006e9ce64e77f3ac165ce2ecf
 size 7014306128