Create README.md

README.md

@@ -0,0 +1,150 @@
+---
+datasets:
+- Muinez/sankaku-webp-256shortest-edge
+---
+
+# StupidAE — d8c16 Tiny Patch Autoencoder
+
+StupidAE is a very small, very fast, and intentionally simple model that still works surprisingly well.  
+It has **13.24M parameters**, compresses by **8× per spatial dimension**, and uses **16 latent channels**.
+
+The main goal: make a AE that doesn’t slow everything down and is fast enough to run directly during text-to-image training.
+
+---
+
+## Code
+
+The code is available on GitHub:
+
+👉 [https://github.com/Muinez/StupidAE](https://github.com/Muinez/StupidAE)
+
+---
+
+## Key Numbers
+
+- Total params: **13,243,539**
+- Compression: **d8 (8×8 patching)**
+- Latent channels: **16 (c16)**
+- Training: **30k steps**, batch size **256**, **~3** RTX 5090-hours
+- Optimizer: **Muon + SnooC**, LR = `1e-3`
+- Trained **without KL loss** (just mse)
+
+---
+
+## Performance (compared to SDXL VAE)
+
+Stats for 1024×1024:
+
+| Component | SDXL VAE | StupidAE |
+|----------|----------|-----------|
+| Encoder FLOPs | 4.34 TFLOPs | **124.18 GFLOPs** |
+| Decoder FLOPs | 9.93 TFLOPs | **318.52 GFLOPs** |
+| Encoder Params | 34.16M | **~3.8M** |
+| Decoder Params | 49.49M | **~9.7M** |
+
+The model is **tens of times faster and lighter**, making it usable directly inside training loops.
+
+---
+
+## Architecture Overview
+
+### ❌ No Attention  
+It is simply unnecessary for this design and only slows things down.
+
+### 🟦 Encoder  
+- Splits the image into **8×8 patches**  
+- Each patch is encoded **independently**  
+- Uses **only 1×1 convolutions**  
+- Extremely fast
+
+The encoder can handle any aspect ratio, but if you want to mix different ARs inside the same batch, the 1×1 conv version becomes inconvenient.
+The Linear encoder version solves this completely — mixed batches work out of the box, although I haven’t released it yet — I can upload it if needed.
+
+There is also a Linear-based encoder version; I can publish it if needed.
+
+### 🟥 Decoder  
+- Uses standard 3×3 convolutions (but 1×1 also works with surprisingly few artifacts)  
+- Uses a **PixNeRF-style head** instead of stacked upsampling blocks
+
+---
+
+## Limitations
+
+- Reconstruction is not perfect — small details may appear slightly blurred.  
+- Current MSE loss: 0.0020.  
+- This can likely be improved by increasing model size.
+
+---
+
+## Notes on 32× Compression
+
+If you want **32× spatial compression**, do **not** use naive 32× patching — quality drops heavily.
+
+A better approach:
+
+1. First stage: patch-8 → 16/32 channels  
+2. Second stage: patch-4 → 256 channels  
+
+This trains much better and works well for text-to-image training too.  
+I’ve tested it, and the results are significantly more stable than naive approaches.
+
+If you want to keep FLOPs low, you could try using patch-16 from the start, but I’m not sure yet how stable the training would be.
+
+I’m currently working on a **d32c64** model with reconstruction quality better than Hunyuan VAE, but I’m limited by compute resources.
+
+---
+
+## Support the Project
+
+I’m renting an **RTX 5090** and running all experiments on it.  
+I’m currently looking for work and would love to join a team doing text-to-image or video model research.
+
+If you want to support development:
+
+- TRC20: 👉 TPssa5ung2MgqbaVr1aeBQEpHC3xfmm1CL  
+- BTC: bc1qfv6pyq5dvs0tths682nhfdnmdwnjvm2av80ej4
+- Boosty: https://boosty.to/muinez  
+
+---
+
+## How to use
+
+Here's a minimal example:
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from torchvision.transforms import v2
+from IPython.display import display
+import requests
+from stae import StupidAE
+
+vae = StupidAE().cuda().half()
+vae.load_state_dict(
+    torch.load(hf_hub_download(repo_id="Muinez/StupidAE", filename="smol_f8c16.pt"))
+)
+
+t = v2.Compose([
+    v2.Resize((1024, 1024)),
+    v2.ToTensor(),
+    v2.Normalize([0.5], [0.5])
+])
+
+image = Image.open(requests.get("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG", stream=True).raw).convert("RGB")
+
+with torch.inference_mode():
+    image = t(image).unsqueeze(0).cuda().half()
+    
+    latents = vae.encode(image)
+    image_decoded = vae.decode(latents)
+    
+    image = v2.ToPILImage()(torch.clamp(image_decoded * 0.5 + 0.5, 0, 1).squeeze(0))
+    display(image)
+```
+---
+## Coming Soon
+
+- Linear-encoder variant  
+- d32c64 model  
+- Tutorial: training text-to-image **without bucketing** (supports mixed aspect ratios)