PrismML-Eng/Bonsai-Image-Demo

Bonsai Image Demo – Quick Overview

What it is – A ready‑to‑run demo for generating images with the Bonsai‑Image diffusion model (4 B parameters). It bundles everything needed to run the model on three common hardware stacks:

Apple Silicon (macOS) – uses the mflux fork with MLX for fast on‑device inference.
Linux with NVIDIA GPU – uses Dropbox’s gemlite backend together with the HQQ 1‑bit/1.58‑bit kernels.
Windows with NVIDIA GPU – same gemlite/HQQ stack via the triton‑windows project, no WSL2 required.

The repo provides scripts to download the appropriate weights, spin up a web‑based Prism Image Studio (FastAPI backend + Next.js frontend), or run a simple CLI for one‑off image generation.

Key Features

Feature	Details
Model variants	`ternary` (1.58‑bit, higher quality) – default; `binary` (1‑bit, smaller & faster).
Cross‑platform backends	macOS → MLX (`mflux`); Linux/Windows → `gemlite` + `HQQ` kernels.
Full studio UI	FastAPI server on `:8000` and Next.js UI on `:3000` (both launched together with `scripts/serve.sh`).
CLI generation	`scripts/generate.sh` for a cold‑start one‑shot run; `scripts/send_request.sh` to talk to a running studio (keeps weights warm).
Automatic dependency handling	`setup.sh` / `setup.ps1` install a pinned Python environment (via `uv`) and Node packages, cloning the required `vendor/` repos.
Supply‑chain safety	Packages newer than `BONSAI_PACKAGE_MIN_AGE_DAYS` (default 7) are rejected to avoid un‑vetted releases.
Model download helper	`scripts/download_model.sh` picks the right weight set for your platform and variant.
Sample sizes	Pre‑defined aspect‑ratio presets for fast (~~0.25 MP) and quality (~~1 MP) renders.

Getting Started (macOS / Linux)

# 1️⃣ Install the environment and download the default ternary model
./setup.sh

# 2️⃣ Start the studio (backend + UI)
./scripts/serve.sh   # backend on 8000, UI on 3000

# 3️⃣ Generate an image via the UI **or** from the terminal
./scripts/send_request.sh -p "A snowy bonsai tree in a misty forest" --size 1248x832

Windows (PowerShell)

# One‑time policy change so scripts can run
Set-ExecutionPolicy -Scope CurrentUser RemoteSigned

# Install & download
.
setup.ps1

# Run the studio
.
scripts
un.ps1   # actually ./scripts/serve.ps1

# Generate
.
scripts
un_request.ps1 -p "A snowy bonsai tree in a misty forest" --size 1248x832

(See scripts/windows.md for driver, vcredist, and OOM troubleshooting.)

Model & Backend Choices

Platform	Default backend	Default model variant
macOS (Apple Silicon)	`mflux` + MLX	`bonsai-image-4B-ternary-mlx`
Linux (NVIDIA)	`gemlite` + HQQ	`bonsai-image-4B-ternary-gemlite`
Windows (NVIDIA)	`gemlite` + HQQ via `triton-windows`	same as Linux

You can override with environment variables, e.g.:

BONSAI_VARIANT=binary ./setup.sh
BONSAI_VARIANT=binary ./scripts/serve.sh

Folder Layout After Setup

Bonsai-image-demo/
  .venv/                     # isolated Python env
  vendor/
    image-studio/            # FastAPI + Next.js code (cloned automatically)
    mflux-prism/             # patched mflux for macOS
  models/
    bonsai-image-4B-ternary-mlx/
    bonsai-image-4B-ternary-gemlite/
    bonsai-image-4B-binary-mlx/      (optional)
    bonsai-image-4B-binary-gemlite/  (optional)
  outputs/                  # generated PNGs
  scripts/
    download_model.*
    serve.*
    generate.*
    send_request.*
    generate.py

Where to Learn More

Website – https://prismml.com
Hugging Face collection – https://huggingface.co/collections/prism-ml/bonsai-image
Whitepaper – bonsai-image-4b-whitepaper.pdf (included in the repo)
Live demo – Hugging Face Space link in the README
Community – Discord server (link in README)

TL;DR

The Bonsai Image Demo is a fully‑scripted, cross‑platform playground for the 4 B‑parameter Bonsai diffusion model. It handles model download, environment setup, and provides both a web UI and a CLI, letting you generate photorealistic images on Apple Silicon or NVIDIA GPUs with minimal fuss.