bonsai-image-ternary-4B-gemlite-2bit-unpacked-encoder

Bonsai Image Ternary 4B — GemLite Transformer with Unpacked Encoder

A Diffusers-compatible variant of Prism ML's Bonsai Image Ternary 4B GemLite release.

This repository preserves the ternary GemLite-packed diffusion transformer, while replacing the original HQQ-packed text encoder with an unpacked, non-HQQ encoder.

The purpose of this variant is to provide a complete pipeline that demonstrates native loading of pre-quantized GemLite checkpoints through the standard Diffusers interface introduced in huggingface/diffusers#14286.

This is primarily a compatibility and integration checkpoint. For the most memory-efficient complete Bonsai Image package, use the original Prism ML release.

What is quantized?

The transformer retains ternary values from the original model and executes its supported linear layers directly through GemLite kernels. The packed transformer weights are not expanded back to full precision during loading.

The text encoder is intentionally unpacked because loading the original HQQ-packed encoder is outside the scope of the current Diffusers GemLite integration.

Why this repository exists

The original Bonsai Image CUDA package combines:

• a GemLite-packed ternary diffusion transformer;
• an HQQ-packed text encoder; and
• the remaining pipeline components.

The Diffusers integration in PR #14286 adds support for loading the GemLite portion of that package through the standard Diffusers quantization interface. It does not add HQQ checkpoint loading.

This variant therefore uses:

• the original GemLite-packed transformer;
• an unpacked encoder that Diffusers can already load; and
• a standardized Diffusers repository layout and quantization configuration.

This allows the complete pipeline to be loaded with:

DiffusionPipeline.from_pretrained(...)

without custom model-loading code.

Status

Native GemLite checkpoint loading is currently proposed in:

• huggingface/diffusers#14286 — Add GemLite quantization backend support for pre-quantized checkpoints

Until that PR is merged and included in a Diffusers release, install Diffusers from the PR author's branch.

The API and installation instructions may change during upstream review.

Requirements

• Python 3.10 or later
• PyTorch with CUDA support
• An NVIDIA GPU supported by PyTorch and Triton
• gemlite>=0.6.0
• Diffusers with GemLite support from PR #14286

GemLite execution is not supported on CPU or Apple MPS.

Installation

While the Diffusers PR is under review

pip install -U \
  "git+https://github.com/gabe-engineers/diffusers.git@add-gemlite-quantizer" \
  "gemlite>=0.6.0" \
  transformers \
  accelerate \
  safetensors

After GemLite support is released in Diffusers

pip install -U \
  diffusers \
  "gemlite>=0.6.0" \
  transformers \
  accelerate \
  safetensors

Usage

import torch
from diffusers import DiffusionPipeline


model_id = (
    "gabe-engineers/"
    "bonsai-image-ternary-4B-gemlite-2bit-unpacked-encoder"
)

pipe = DiffusionPipeline.from_pretrained(
    model_id,
    dtype=torch.float16,
    device_map="cuda",
)

image = pipe(
    prompt="A bonsai tree in a quiet ceramic studio, soft morning light",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=1.0,
).images[0]

image.save("bonsai-gemlite.png")

No explicit GemLiteConfig is required when loading this repository. The serialized GemLite quantization configuration is stored with the transformer checkpoint and is detected automatically by Diffusers.

Dtype requirement

The dtype passed to from_pretrained must match the compute_dtype stored in the GemLite quantization configuration.

This checkpoint uses:

dtype=torch.float16

Passing a different dtype will raise an error rather than silently loading the pipeline with an incompatible configuration.

Differences from the original release

This repository is derived from:

• prism-ml/bonsai-image-ternary-4B-gemlite-2bit

The main difference is the text encoder:

Because the encoder is unpacked, this repository has a larger total storage and memory footprint than the fully compressed original release.

The transformer itself remains GemLite-packed and retains the primary purpose of this checkpoint: validating native low-bit GemLite loading in Diffusers.

Intended uses

This repository is intended for:

• testing native GemLite checkpoint loading in Diffusers;
• reproducing the documentation example from PR #14286;
• experimenting with ternary image-generation transformers on NVIDIA GPUs;
• validating saving and reloading of GemLite-backed Diffusers models;
• serving as a reference checkpoint for future GemLite integrations.

Limitations

• Only pre-quantized GemLite checkpoints are supported. The Diffusers integration does not quantize an unquantized model into GemLite format at load time.
• The unpacked encoder uses more disk and device memory than the HQQ-packed encoder in the original release.
• CPU execution is not supported because GemLite relies on Triton GPU kernels.
• Apple Silicon and MPS are not supported by this CUDA checkpoint.
• Disk offloading is not supported by the current GemLite quantizer.
• Parallel checkpoint loading is disabled for this backend.
• This repository should not be used to evaluate the total memory footprint of the original fully compressed Prism ML deployment.
• Upstream Diffusers support is still under review and may change before release.

Model provenance

The model architecture and trained weights originate from the Bonsai Image releases by Prism ML.

Relevant upstream repositories:

• Bonsai Image Ternary 4B — GemLite 2-bit
• Bonsai Image Ternary 4B — Unpacked
• Bonsai Image collection

This repository does not claim authorship of the underlying Bonsai Image model. Its changes are limited to checkpoint composition, repository layout, serialized quantization metadata, and compatibility with the proposed Diffusers GemLite loader.

GemLite

GemLite provides Triton kernels and packed representations for low-bit matrix operations.

The Diffusers integration replaces eligible torch.nn.Linear modules with GemLite linear modules before loading the checkpoint. The packed weights, scales, zero points, metadata, and original tensor shapes are then restored directly from the serialized checkpoint.

Diffusers integration

Implementation and discussion:

• huggingface/diffusers#14286

The integration is scoped to:

• detecting serialized GemLite quantization configurations;
• replacing eligible linear layers with GemLite modules;
• loading packed GemLite state through the low-memory Diffusers loader;
• supporting device-map memory estimation using packed tensor sizes;
• running inference without dequantizing the transformer weights;
• saving and reloading GemLite-backed checkpoints.

HQQ encoder support is intentionally treated as a separate integration concern.

License

This repository is distributed under the Apache License 2.0.

Refer to the included LICENSE and NOTICE.md files and the upstream model repositories for complete attribution and licensing information.