Rewrite model card

README.md

@@ -1,90 +1,149 @@
 ---
 license: mit
+language:
+  - en
 library_name: pytorch
 pipeline_tag: text-to-video
 tags:
   - text-to-video
+  - video-generation
   - diffusion-transformer
   - rectified-flow
   - stickers
-  - from-scratch
-  - tiny
+  - animation
 ---
 
+<div align="center">
+
 # Bhaav-E
 
-A **212M-parameter text-to-video model** that makes **2-second animated stickers** (256 × 256, 17 frames, 8 fps, looping).
-It was trained from scratch for about $12 of GPU time on 75K stickers generated by code.
+**A compact text-to-video model for animated stickers**
 
-**Try it free:** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Mohit-5899/bhaav-e/blob/main/bhaav_e_demo.ipynb) (runs the web demo on a free Colab GPU) · [showcase page](https://huggingface.co/spaces/Mohit5899/bhaav-e)
+[![Project Page](https://img.shields.io/badge/Project-Page-e8590c)](https://huggingface.co/spaces/Mohit5899/bhaav-e)
+[![GitHub](https://img.shields.io/badge/GitHub-Code-181717?logo=github)](https://github.com/Mohit-5899/bhaav-e)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Mohit-5899/bhaav-e/blob/main/bhaav_e_demo.ipynb)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/Mohit-5899/bhaav-e/blob/main/LICENSE)
 
-Code, data generator, training logs and full evaluation: **[github.com/Mohit-5899/bhaav-e](https://github.com/Mohit-5899/bhaav-e)**
+</div>
 
 | | | | |
-|---|---|---|---|
+|:---:|:---:|:---:|:---:|
 | ![](examples/00_seed0.gif) | ![](examples/01_seed0.gif) | ![](examples/02_seed0.gif) | ![](examples/03_seed0.gif) |
 | *a crying emoji bouncing, tears falling, light blue background* | *a red furious emoji pulsing, steam puffing from its head, dark navy background* | *a blue success checkmark drawing itself inside a circle, confetti bursting out, white background* | *a party popper emoji pulsing, sparkles twinkling around it, light grey background* |
 | ![](examples/04_seed0.gif) | ![](examples/05_seed0.gif) | ![](examples/06_seed0.gif) | ![](examples/07_seed0.gif) |
 | *a cool emoji with sunglasses floating gently, sparkles twinkling around it, peach background* | *a glowing star emoji wobbling side to side, sparkles twinkling around it, dark navy background* | *a purple loading spinner spinning round and round, pulsing, black background* | *a screaming in shock emoji wobbling side to side, peach background* |
 
-Seed 0, 50 steps, guidance 3, first sample for each prompt (not cherry-picked).
+<sub>Uncurated samples (seed 0, 50 steps, guidance 3.0): the first and only sample generated for each prompt.</sub>
+
+## Model Details
+
+Bhaav-E generates 2-second looping stickers from a text prompt. It is a 211.92M-parameter video diffusion
+transformer trained from scratch with rectified flow, operating in the latent space of the frozen Wan2.2 video VAE and
+conditioned on a frozen Flan-T5-Base text encoder.
+
+- **Developed by:** Mohit
+- **Model type:** latent video diffusion transformer (rectified flow)
+- **Output:** 17 frames, 256 × 256, 8 fps, looping
+- **Language:** English prompts
+- **License:** MIT (frozen components: Apache-2.0)
+- **Code:** [github.com/Mohit-5899/bhaav-e](https://github.com/Mohit-5899/bhaav-e)
+
+| Component | Specification |
+|---|---|
+| Denoiser (trained) | 16 transformer blocks, width 768, 12 heads, 2 × 2 patches (320 tokens) |
+| Text encoder (frozen) | [google/flan-t5-base](https://huggingface.co/google/flan-t5-base), max 35 tokens |
+| Autoencoder (frozen) | [Wan2.2 VAE](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers), 48 latent channels |
+| Sampler | 50 Euler steps, classifier-free guidance 3.0, time shift 5.0 |
 
 ## Files
 
-| file | what |
+| File | Description |
 |---|---|
-| `bhaav_e.safetensors` | r2 (release) VideoDiT, EMA weights, fp32, step 40,000 |
-| `config.json` | model size, sampler defaults |
-| `text/info.json`, `latents/vae_config.json` | caption-encoder settings and VAE latent statistics, read by `sample.py` |
-| `r1/` | the first run (natural data mix), same format, for comparison |
-| `code/videodit.py`, `code/sample.py` | the model definition and the sampler |
+| `bhaav_e.safetensors` | Bhaav-E (recommended): EMA weights, step 40,000, fp32 |
+| `config.json` | model size and sampler defaults |
+| `text/info.json`, `latents/vae_config.json` | text-encoder settings and VAE latent statistics used at inference |
+| `r1/` | baseline checkpoint trained on the natural attribute distribution, same format |
+| `code/` | model definition (`videodit.py`) and sampler (`sample.py`) |
+
+## Uses
 
-## Model
+**Intended use:** generating short animated stickers, reaction emoji and UI micro-animations, and research on small
+video diffusion models, compositional generalization and evaluation with procedurally generated data.
 
-- **VideoDiT** (trained, ours): 16 transformer blocks, width 768, 12 heads, 2×2 patches over a 48 × 5 × 16 × 16 latent
-  (320 tokens). Caption through cross-attention, noise level through adaptive LayerNorm.
-- **Frozen, borrowed:** [Wan2.2 VAE](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers) (video ↔ latent),
-  [Flan-T5-Base](https://huggingface.co/google/flan-t5-base) encoder (caption → vectors).
-- **Training:** rectified flow (predict `v = noise − clean`), Muon + AdamW, batch 64 × 40,000 steps, bf16, one RTX PRO
-  4500 for about 3.6 h (≈ $2.60). Captions dropped 10% of the time for classifier-free guidance.
+**Out of scope:** subjects outside the training vocabulary, photorealistic video, high resolution or long clips.
 
-## Usage
+## How to Use
 
-Needs a CUDA GPU (tested on a 24 GB L4).
+**No installation:** open the [Colab notebook](https://colab.research.google.com/github/Mohit-5899/bhaav-e/blob/main/bhaav_e_demo.ipynb),
+select **Runtime → Run all**, and open the `gradio.live` link it prints. This runs the full web demo on a free GPU.
+
+**Python** (CPU or GPU):
 
 ```bash
+git clone https://github.com/Mohit-5899/bhaav-e && cd bhaav-e
 pip install torch diffusers transformers sentencepiece safetensors pillow numpy huggingface_hub
-hf download Mohit5899/bhaav-e --local-dir bhaav-e
-cd bhaav-e/code && python sample.py --ckpt ../bhaav_e.safetensors --cache .. \
-    --prompt "a crying emoji bouncing, tears falling, light blue background" --out ../samples
 ```
 
-Write prompts in the training style: *subject, motion, effect, background*. Known subjects are 10 emotion faces
-(laugh, crying, heart-eyes, facepalm, angry, shocked, wink, thinking, sleepy, cool), 10 animated emoji (fire, heart,
-party popper, thumbs-up, star, 100, clap, rocket, trophy, light bulb) and 10 UI animations (spinner, dots loader,
-success, error, like, bell, download, progress, toggle, typing). Effects: tears, tears of joy, hearts, sparkles,
-sweat, steam, zzz, thought bubble, confetti, embers, exhaust, none. There are 12 plain background colours.
+```python
+import sys; sys.path.insert(0, "space")
+from bhaav import Bhaav, save_gif
+
+pipe = Bhaav()                       # downloads this checkpoint, the VAE and the text encoder
+frames = pipe.generate("a crying emoji bouncing, tears falling, light blue background",
+                       steps=50, cfg=3.0, seed=0)
+save_gif(frames, "sticker.gif")
+```
+
+**Prompt format:** *subject, motion, effect, background*. The model supports 30 subjects (10 emotion faces, 10
+animated emoji, 10 UI animations), 9 motions, 12 effects and 12 background colours; the full list is in the
+[GitHub README](https://github.com/Mohit-5899/bhaav-e#prompt-format).
+
+## Training
+
+| | |
+|---|---|
+| Data | 85,816 procedurally rendered clips; 61,571 used for training, with effects balanced per subject |
+| Objective | rectified flow, MSE on predicted velocity |
+| Optimizer | Muon (lr 0.02) for transformer weight matrices, AdamW (lr 2e-4) for other parameters |
+| Batch size × steps | 64 × 40,000 |
+| Precision | bf16 autocast |
+| Compute | 1 × NVIDIA RTX PRO 4500 (32 GB), ≈ 3.6 hours |
 
 ## Evaluation
 
-There are 256 test prompts with combinations seen in training, and 256 with combinations held out from training. Each
-generated sticker is checked automatically for subject, motion, effect, background and accent. A clip counts as a model
-error only if the checker reads the real version of that clip correctly.
+256 prompts with attribute combinations seen in training and 256 with combinations held out from training. Automatic
+checkers read subject, motion, effect, background and accent from each generated clip; errors are counted only where
+the checker reads the corresponding real clip correctly.
 
-| | seen | unseen |
+| Model | Seen combinations | Unseen combinations |
 |---|---|---|
-| r1 (natural mix) | 2.0% | 5.3% |
-| **r2 (this model, effects balanced per subject)** | **2.0%** | **3.3%** |
+| **Bhaav-E** | **2.0%** | **3.3%** |
+| Bhaav-E r1 (baseline) | 2.0% | 5.3% |
 
-The unseen difference is not statistically significant (one seed each, sign test p = 0.23). The main remaining
-failure is adding a subject's usual effect when none was asked for.
+Error rate on any attribute, guidance 3.0. The dominant failure mode is effect leakage: adding the effect most
+associated with a subject when none is requested. The unseen difference between the models is not statistically
+significant (exact sign test, p = 0.23; one training seed each). Full report:
+[runs/r2/RESULTS.md](https://github.com/Mohit-5899/bhaav-e/blob/main/runs/r2/RESULTS.md).
 
 ## Limitations
 
-A closed, tiny world: 30 subjects, 12 effects, 12 backgrounds, 256 px. It does not draw subjects outside the training
-set. Image-to-sticker (given first frame) was trained but the sampler does not expose it yet.
+- Closed vocabulary: only the 30 trained subjects, 12 effects and 12 backgrounds.
+- 256 × 256 resolution and 2-second duration.
+- Unrequested effects are occasionally added.
+- Evaluated with one seed per model and about 250 prompts per split.
+
+## Citation
+
+```bibtex
+@misc{bhaave2026,
+  title        = {Bhaav-E: A Compact Text-to-Video Model for Animated Stickers},
+  author       = {Mohit},
+  year         = {2026},
+  howpublished = {\url{https://github.com/Mohit-5899/bhaav-e}}
+}
+```
 
-## Credits
+## Acknowledgements
 
-The emoji art in the training data is [Microsoft Fluent Emoji](https://github.com/microsoft/fluentui-emoji) (MIT). The
-Wan2.2 VAE and Flan-T5 are Apache-2.0.
+Built on [Wan2.2](https://github.com/Wan-Video/Wan2.2), [Flan-T5](https://huggingface.co/google/flan-t5-base) and
+[Microsoft Fluent Emoji](https://github.com/microsoft/fluentui-emoji) (MIT, training artwork).

code/videodit.py

@@ -4,7 +4,7 @@ Design source: the DiT idea (Peebles & Xie, 2023) at the published "s210" size (
 12 heads, about 212M parameters). We copy the blueprint, not code; `count_params` checks we built the same size.
 
 What it does, in one line: given a noisy sketch x_t, the noise level t and a caption, predict the velocity
-v = noise - clean (rectified flow, see how-training-works.md 5a).
+v = noise - clean (rectified flow, see docs/how-training-works.md 5a).
 
 Shapes (B = batch):
   x      [B, 48, 5, 16, 16]   VAE latent: 48 channels, 5 time steps, 16 x 16 positions