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?
Component
Format
Diffusion transformer
Ternary weights packed in GemLite's 2-bit representation
Text encoder
Unpacked, non-HQQ weights
VAE
Standard pipeline weights
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:
1import torch
2from diffusers import DiffusionPipeline
345model_id =(6"gabe-engineers/"7"bonsai-image-ternary-4B-gemlite-2bit-unpacked-encoder"8)910pipe = DiffusionPipeline.from_pretrained(11 model_id,12 dtype=torch.float16,13 device_map="cuda",14)1516image = pipe(17 prompt="A bonsai tree in a quiet ceramic studio, soft morning light",18 height=1024,19 width=1024,20 num_inference_steps=4,21 guidance_scale=1.0,22).images[0]2324image.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.
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.
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.
