# Bonsai Image 4B — Local Diffusion Model Setup

## Overview

PrismML's Bonsai Image 4B is a family of compact diffusion transformers for text-to-image, based on FLUX.2 Klein 4B architecture. Available in ternary (1.58-bit, higher quality) and binary (1-bit, smaller).

## Variants & Sizes (NVIDIA CUDA / gemlite)

| Variant | Transformer | Total Payload | Quality vs FP16 | Speed (1024²) |
|---------|------------|---------------|-----------------|---------------|
| Ternary (1.58-bit) | 1.21 GB | 4.55 GB | 95% | ~4.5s RTX 3080 |
| Binary (1-bit) | 0.93 GB | 4.09 GB | 88% | ~3.5s RTX 3080 |

Peak VRAM at 1024²: ~6.8 GiB on RTX 3080 — fits 11 GB 2080 Ti with headroom.

## Setup

```bash
git clone https://github.com/PrismML-Eng/Bonsai-Image-Demo
cd Bonsai-Image-Demo
./setup.sh
```

This installs Python deps (via `uv`), clones vendor repos (image-studio, mflux-prism), and downloads model weights to `models/`.

### CUDA / Driver Compatibility Fix

The `setup.sh` uv.lock resolves to torch 2.12.0+cu130, which requires NVIDIA driver >= 570.x (CUDA 13.0). On systems with older drivers (e.g., 550.163.01 / CUDA 12.4), torch fails with `The NVIDIA driver on your system is too old`.

**Fix** — pin torch + gemlite after `setup.sh`:

```bash
source .venv/bin/activate
uv pip install 'torch==2.6.0+cu124' --index-url https://download.pytorch.org/whl/cu124
uv pip install 'gemlite==0.4.7'
```

This gives torch 2.6.0+cu124 + gemlite 0.4.7, compatible with driver 550.x. The `float8_e8m0fnu` attribute (needed by gemlite ≥0.5.x) only exists in torch ≥2.12+cu130 — gemlite 0.4.7 doesn't reference it.

**Alternative:** Upgrade NVIDIA driver to ≥570.x and re-run `setup.sh` with defaults.

## Running

**Full studio** (FastAPI :8000 + Next.js :3000):
```bash
./scripts/serve.sh
```

**One-shot CLI** (cold-start each call):
```bash
./scripts/generate.sh -p "your prompt" --size 1024x1024 --seed 42
```

**HTTP client to running studio** (weights stay warm):
```bash
./scripts/send_request.sh -p "your prompt" --size 1248x832
```

## Configuration

| Variable | Default | Description |
|----------|---------|-------------|
| `BONSAI_VARIANT` | `ternary` | `ternary` or `binary` |
| `BACKEND_PORT` | 8000 | FastAPI port |
| `FRONTEND_PORT` | 3000 | Next.js port |

## Model Details

- **Sampler:** FlowMatchEuler-discrete, 4 steps, guidance=1.0, shift=3.0
- **No CFG, no negative prompts** — natural language prompting only
- **Text encoder:** Qwen3-4B at 4-bit HQQ (offloaded after prompt encode)
- **VAE:** Flux2 32-channel latent, tiled decode (128px tiles)
- **Native resolution:** 1024×1024, any multiple of 32 works
- **Suggested sizes:** 512×512 (fast), 1024×1024 (quality), 1248×832 (landscape), 832×1248 (portrait)

## Software Stack

- **gemlite** — fused low-bit GEMM for NVIDIA (https://github.com/mobiusml/gemlite)
- **HQQ** — low-bit quantization runtime (https://github.com/mobiusml/hqq)
- **diffusers** — HuggingFace pipeline (Flux2KleinPipeline)
- **4-step threshold:** running more steps doesn't improve quality and can introduce artifacts

## Resources

- **Demo repo:** https://github.com/PrismML-Eng/Bonsai-Image-Demo
- **White paper:** `bonsai-image-4b-whitepaper.pdf` in demo repo
- **HuggingFace:** prism-ml/bonsai-image-ternary-4B-gemlite-2bit
- **Discord:** https://discord.gg/prismml