Blame for README.md

# MOSS-Image — EntroPack Pre-Quantized Models

We provide pre-quantized weight packages produced by [**EntroPack**](https://github.com/modelscope/entropack) for **MOSS-Image** (Base and SFT, from [openmoss/MOSS-Image-Base](https://www.modelscope.cn/models/openmoss/MOSS-Image-Base) and [openmoss/MOSS-Image-SFT](https://www.modelscope.cn/models/openmoss/MOSS-Image-SFT)), which can be used with the quantization features of [**DiffSynth-Studio**](https://github.com/modelscope/DiffSynth-Studio). Use them with ComfyUI through [DiffSynth-ComfyUI](https://github.com/modelscope/DiffSynth-ComfyUI).

The packages run with as little as 4 GB of VRAM.

Base is provided at four levels, SFT additionally ships a 3.5-bit variant. Base and SFT share the same text encoder packages.

**BF16 GEMM levels** (the 3.5-bit and 4.5-bit DiTs use mixed quantization):

* **`3.5bit`** (SFT only, lowest VRAM) — mixed-bit: the quantization-sensitive layers keep **5.0 bpp**, while the remaining layers use **3.2516 bpp**, giving **3.526 bpp** on the package;
* **`4.5bit`** (balanced) — mixed-bit: the quantization-sensitive layers keep **5.0 bpp**, while the remaining layers use **4.4172 bpp**, giving **4.545 bpp** on the package;
* **`6bit`** (highest precision of the BF16 GEMM levels) — uniform entropy coding, **6.055 bpp** on the package.

**INT8 GEMM levels** use a hybrid allocation:

* **`int8_6bit`** (fastest inference) — the quantization-sensitive linears including embeddings, time embeddings, modulation and final layers are entropy-coded at **8.5 bits per weight** and run on **BF16 GEMM**, while all remaining linears are **INT8 W8A8-quantized with a ConvRot rotation** and run on **INT8 GEMM**; the package averages **5.999 bpp**;
* **`int8_7bit`** (highest precision of the INT8 levels) — the same hybrid allocation, averaging **7.004 bpp** on the package.

The INT8 GEMM path requires Triton and a CUDA GPU with compute capability 8.0 or newer.

All packages are entropy-coded from the original BF16 weights directly. Every `*.safetensors` ships with a same-named `*.json` that records its quantization config. `actual bpp = 8 × Σ storage bytes / Σ numel` is measured on the package artifacts.

## Preview

All previews are generated with the same prompt and seed; `bf16` is the full-precision reference. The quantized previews are ordered by PSNR against the `bf16` reference (highest first).

### SFT

<table width="100%">
  <tr>
    <td width="33%" align="center">bf16 (reference)</td>
    <td width="33%" align="center">int8@7</td>
    <td width="33%" align="center">int8@6</td>
  </tr>
  <tr>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/bf16.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/dit_int8_7bit.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/dit_int8_6bit.jpg"></td>
  </tr>
  <tr>
    <td width="33%" align="center">6bit</td>
    <td width="33%" align="center">4.5bit</td>
    <td width="33%" align="center">3.5bit</td>
  </tr>
  <tr>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/dit_6bit.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/dit_4.5bit.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-SFT/dit_3.5bit.jpg"></td>
  </tr>
</table>

### Base

<table width="100%">
  <tr>
    <td width="33%" align="center">bf16 (reference)</td>
    <td width="33%" align="center">6bit</td>
    <td width="33%" align="center">int8@7</td>
  </tr>
  <tr>
    <td align="center"><img width="100%" src="assets/MOSS-Image-Base/bf16.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-Base/dit_6bit.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-Base/dit_int8_7bit.jpg"></td>
  </tr>
  <tr>
    <td width="33%" align="center">int8@6</td>
    <td width="33%" align="center">4.5bit</td>
    <td></td>
  </tr>
  <tr>
    <td align="center"><img width="100%" src="assets/MOSS-Image-Base/dit_int8_6bit.jpg"></td>
    <td align="center"><img width="100%" src="assets/MOSS-Image-Base/dit_4.5bit.jpg"></td>
    <td></td>
  </tr>
</table>

## DiT

### SFT

| name | file | actual bpp | size |
|---|---|---|---|
| dit_3.5bit (mixed) | `models/MOSS-Image-SFT/dit_3.5bit.safetensors` | 3.526 | 2.24 GB |
| dit_4.5bit (mixed) | `models/MOSS-Image-SFT/dit_4.5bit.safetensors` | 4.545 | 2.89 GB |
| dit_6bit | `models/MOSS-Image-SFT/dit_6bit.safetensors` | 6.055 | 3.85 GB |
| dit_int8_6bit (hybrid INT8 GEMM + ConvRot) | `models/MOSS-Image-SFT/dit_int8_6bit.safetensors` | 5.999 | 3.81 GB |
| dit_int8_7bit (hybrid INT8 GEMM + ConvRot) | `models/MOSS-Image-SFT/dit_int8_7bit.safetensors` | 7.004 | 4.45 GB |

### Base

| name | file | actual bpp | size |
|---|---|---|---|
| dit_4.5bit (mixed) | `models/MOSS-Image-Base/dit_4.5bit.safetensors` | 4.545 | 2.89 GB |
| dit_6bit | `models/MOSS-Image-Base/dit_6bit.safetensors` | 6.055 | 3.85 GB |
| dit_int8_6bit (hybrid INT8 GEMM + ConvRot) | `models/MOSS-Image-Base/dit_int8_6bit.safetensors` | 5.999 | 3.81 GB |
| dit_int8_7bit (hybrid INT8 GEMM + ConvRot) | `models/MOSS-Image-Base/dit_int8_7bit.safetensors` | 7.004 | 4.45 GB |

## Text Encoder

The text encoder packages are shared by Base and SFT.

| name | file | actual bpp | size |
|---|---|---|---|
| text_encoder_4.5bit | `models/text_encoder_4.5bit.safetensors` | 4.552 | 5.97 GB |
| text_encoder_6bit | `models/text_encoder_6bit.safetensors` | 6.070 | 7.29 GB |
| text_encoder_int8_6bit (hybrid INT8 GEMM + ConvRot) | `models/text_encoder_int8_6bit.safetensors` | 6.000 | 7.23 GB |
| text_encoder_int8_7bit (hybrid INT8 GEMM + ConvRot) | `models/text_encoder_int8_7bit.safetensors` | 6.993 | 8.08 GB |

## Usage

### Install

```bash
# DiffSynth-Studio
git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio && pip install -e .

# EntroPack (pick the extra matching your CUDA version: cuda13 / cuda12)
git clone https://github.com/modelscope/entropack.git
cd entropack && pip install -e '.[cuda13]'
```

### Example scripts

* Inference: [standard](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/moss_image/model_inference/MOSS-Image-EntroPack.py) / [low VRAM](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/moss_image/model_inference_low_vram/MOSS-Image-EntroPack.py) — defaults to SFT INT8 6bit; switch `dit_package` / `base_model_id` together to use Base or another level
* LoRA training: [train](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/moss_image/model_training/lora/MOSS-Image-entropack.sh) / [validation](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/moss_image/model_training/validate_lora/MOSS-Image-entropack.py) — defaults to SFT; replace `MOSS-Image-SFT` with `MOSS-Image-Base` in the dataset and model paths to train on Base

`dit_package` and `text_encoder_package` are chosen independently: any combination works (e.g. `dit_int8_6bit` with `text_encoder_4.5bit`). The 6-bit packages are pre-registered in DiffSynth-Studio, so they load with a plain `ModelConfig`; the helper below registers any package's json sidecar before loading (the INT8 packages also read their ConvRot settings and 8.5-bit sensitive-layer allocation from it):

```python
import json

import torch

from diffsynth.configs import MODEL_CONFIGS
from diffsynth.pipelines.moss_image import MOSSImagePipeline, ModelConfig

prequant_model_id = "DiffSynth-Studio/MOSS-Image-EntroPack"
base_model_id = "openmoss/MOSS-Image-SFT"
dit_package = "models/MOSS-Image-SFT/dit_4.5bit.safetensors"
text_encoder_package = "models/text_encoder_4.5bit.safetensors"


def make_package(pattern):
    package = ModelConfig(model_id=prequant_model_id, origin_file_pattern=pattern)
    package.download_if_necessary()
    sidecar = ModelConfig(model_id=prequant_model_id, origin_file_pattern=pattern.replace(".safetensors", ".json"))
    sidecar.download_if_necessary()
    with open(sidecar.path) as handle:
        entry = json.load(handle)
    if not any(config.get("model_hash") == entry["model_hash"] for config in MODEL_CONFIGS):
        MODEL_CONFIGS.append(entry)
    return package


pipe = MOSSImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        make_package(dit_package),
        make_package(text_encoder_package),
        ModelConfig(model_id=base_model_id, origin_file_pattern="vae/diffusion_pytorch_model.safetensors", computation_dtype=torch.float32),
    ],
    tokenizer_config=ModelConfig(model_id=base_model_id, origin_file_pattern="tokenizer/"),
)

image = pipe(
    prompt="""Full-color cinematic outdoor night photograph of a stunning young Asian woman in her early twenties with a natural warm ivory complexion and a healthy glow, wearing a soft blush pink off-shoulder dress with a fitted bodice, framed from the waist up against softly illuminated historic architecture. She stands near the centre of the frame, facing the camera with a gentle closed-lip smile, her hands resting loosely clasped in front of her waist with relaxed elbows; she has a slim oval face with softly defined cheekbones and a delicate narrow jaw, dark expressive eyes, natural-shaped brows, and subtle natural makeup with soft red lips. Her long black hair falls in loose natural waves over her shoulders, with fine strands framing her face. The fabric of the dress falls in soft folds across her shoulders and chest. In the frame-left background, the word "EntroPack" glows as a hand-drawn bubbly yellow neon sign mounted on a softly floodlit dark stone building above the street; rows of warm-lit windows, stone cornices and shopfronts remain clearly readable around the sign in gentle amber light, while the thick rounded irregular letters ride a gently bouncing baseline with a bright core and a soft outer glow, kept sharp and unblurred against the night. The background recedes into deep blue tones: the Eiffel Tower rises above distant buildings on the frame-right as a softly silhouetted tapering lattice outline in the shallow depth of field, while small amber, muted turquoise and rose-coloured lights dissolve into rounded bokeh across the upper background, leaving a darker area behind her head. Soft neutral frontal illumination with gentle falloff reveals her face with natural warm skin tone and subtle texture, a faint cooler rim separates her hair from the night, and the skirt of her dress fills the lower edge of the frame.""",
    height=1024, width=1024,
    seed=123, num_inference_steps=50, cfg_scale=3.0,
)
image.save("image.jpg")
```