SumomoLee's picture
v2: core-compatible text encoder keys (model.*), core workflows, Comfy-Org comparison, VAE description fix
16364ff verified
|
Raw History Blame Contribute Delete
14.7 kB
---
license: apache-2.0
base_model: OpenMOSS-Team/MOSS-Image-Base
base_model_relation: quantized
pipeline_tag: text-to-image
library_name: comfyui
tags:
- comfyui
- moss-image
- text-to-image
- quantized
- int8
- int4
- convrot
- hadamard
- diffusion-single-file
---
# MOSS-Image-Base ConvRot INT8 / INT4 for ComfyUI
ComfyUI-native **INT8 ConvRot** and **INT4 ConvRot (W4A4)** quantizations of
[OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base), the base text-to-image model of the MOSS-Image family by the
[OpenMOSS Team](https://github.com/OpenMOSS/MOSS-Image).
These files use the same layouts, tensor names and quantization metadata as the community SFT quantizations below, so they load with the same loaders:
- INT4: [toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI](https://huggingface.co/toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI)
- INT8: [Abiray/MOSS-Image-SFT-INT8-convrot](https://huggingface.co/Abiray/MOSS-Image-SFT-INT8-convrot)
Both reference repos target **MOSS-Image-SFT**. This repo applies the identical recipe to **MOSS-Image-Base**, the non-post-trained base model,
which is the better starting point for LoRA training. As a check, our converter applied to SFT weights reproduces both reference files bit-exactly.
## Files
| File | Variant | Size |
| :-- | :-- | --: |
| `diffusion_models/moss-image-base-dit-int8-convrot.safetensors` | DiT, INT8 ConvRot (W8, per-channel) | 5.41 GB (5,413,723,800 bytes) |
| `diffusion_models/moss-image-base-dit-int4-convrot.safetensors` | DiT, INT4 ConvRot W4A4 | 3.04 GB (3,035,537,344 bytes) |
| `text_encoders/moss-image-qwen35-int8-convrot.safetensors` | Qwen3.5 text encoder, INT8 ConvRot (optional) | 8.45 GB (8,452,959,908 bytes) |
| `vae/moss-image-vae.safetensors` | official MOSS-Image VAE (FLUX.2 latent space, slim decoder), unchanged FP32 copy | 0.25 GB (249,521,372 bytes) |
| `workflows/core/*.json` | workflows for native ComfyUI core support (stock nodes), UI and API format | |
| `workflows/*.json` | workflows for the MOSS-Image plugin, UI and API format | |
| `quantize_moss.py` | conversion script | |
The BF16 DiT is ~10.2 GB and the BF16 text encoder ~16.8 GB.
MOSS-Image-Base and MOSS-Image-SFT ship byte-identical text encoder and VAE weights, so the text encoder here is identical to the one in the Abiray INT8 repo.
The text encoder uses HF / ComfyUI-core key names (`model.language_model.*`, `model.visual.*`). Both ComfyUI core and the MOSS-Image plugin accept them.
You can also keep the BF16 text encoder, exported with the plugin's `export.py`.
## Quantization method
**ConvRot** rotates each linear layer with a block-diagonal, normalized regular Hadamard matrix over groups of the input dimension before quantizing.
This spreads per-channel outliers across the group, so plain symmetric absmax quantization works well.
Quantization uses [comfy-kitchen](https://github.com/Comfy-Org/comfy-kitchen) 0.2.37 layouts, the same ones ComfyUI uses natively, applied to the FP32-upcast BF16 weights.
**INT8 (`int8_tensorwise` + `convrot`)**, same as Abiray/MOSS-Image-SFT-INT8-convrot
- Hadamard rotation group 256, int8 symmetric weights, one FP32 scale per output channel (`weight_scale` [N,1]).
- Per-layer config is stored as a `<layer>.comfy_quant` JSON tensor.
- At runtime ComfyUI uses an INT8 GEMM on GPUs that support it, otherwise it dequantizes.
- Mean relative weight error: DiT 0.868%, text encoder 0.868%.
**INT4 (`convrot_w4a4`)**, same as toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI
- Hadamard rotation group 256, signed int4 weights packed two per byte, one FP32 scale per output channel, kernel group size 64.
- Config is stored in the safetensors header `_quantization_metadata`.
- At runtime activations are rotated online and quantized to int4 per token, and the matmul runs on int4 tensor cores (**W4A4**).
- Mean relative weight error: DiT 15.74%.
**Quantized layers (DiT, both variants):** 252 linears, i.e. attention `proj.q`, `proj.k`, `proj.v`, `proj.o` and MLP `proj.ffn_up_gate` and `proj.ffn_down` in all 42 blocks.
**Kept in original precision (DiT):** text embedder (`text_embedder.*`), time embedding, `shared_block_modulation`, `video_head_modulation`,
input embedder (`video_embedder`), output projection (`final_linear_video`), all RMSNorm / q-norm / k-norm weights, and `rope.bands`
(FP32 in MOSS-Image-Base, kept as FP32).
**Text encoder (INT8 only):** 309 tensors are quantized: `embed_tokens` (per-row, rotated), all language-model attention / linear-attention / MLP projections
except `linear_attn.in_proj_a` and `in_proj_b`, and the vision-tower block linears. The rotation group is the largest power of 4 (at most 256) that divides
the input dimension: 256, 64, or 16 for `visual.*.linear_fc2`. Norms, biases, `A_log`, `dt_bias`, `conv1d`, `in_proj_a/b`, the vision merger, patch embed and
position embedding stay BF16.
## Validation
**GPU image-quality and speed benchmarks were not run for this release.** There are no PSNR / SSIM / LPIPS / CLIP numbers,
VRAM measurements or s/it figures here. Please evaluate on your own prompts and hardware, and compare against the BF16 model, before relying on these files.
INT4 W4A4 also quantizes activations, so expect visible differences from BF16, larger than with INT8.
What was checked:
**Per-layer weight reconstruction error** (relative Frobenius error of the dequantized weight vs the BF16 weight, over all quantized layers):
| File | Quantized layers | Mean rel. error | Max rel. error |
| :-- | --: | --: | --: |
| INT8 ConvRot DiT | 252 | 0.868% | 0.980% |
| INT8 ConvRot text encoder | 309 | 0.868% | 1.020% |
| INT4 ConvRot W4A4 DiT | 252 | 15.7% | 17.7% |
This measures weight error only. It does not include INT4 activation quantization or the effect on generated images.
**Format fidelity:** the same converter (`quantize_moss.py`), applied to MOSS-Image-SFT layers, reproduces the quantized weights and scales of
[toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI](https://huggingface.co/toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI) and
[Abiray/MOSS-Image-SFT-INT8-convrot](https://huggingface.co/Abiray/MOSS-Image-SFT-INT8-convrot) bit-exactly on the layers tested.
The DiT files have the same tensor names, shapes and dtypes and the same quantization configs as those repos. The one difference is `rope.bands`,
which is FP32 in MOSS-Image-Base and is kept as FP32. The INT8 text encoder has the same keys, shapes and dtypes as Abiray's, and the tensors sampled are byte-identical.
**ComfyUI load / generate test (CPU only):** in headless ComfyUI `bdc63e7` with the official MOSS-Image plugin `1f808b7`:
- Both DiT files and the INT8 text encoder load through `MOSSImageUNETLoader` / `MOSSImageCLIPLoader`.
- The files appear in the loader lists.
- An API-format prompt (INT4 DiT + INT8 TE + VAE, low resolution, 2 steps) ran end to end and saved an image. The INT8 DiT generated an image the same way.
- The same test passes on ComfyUI core with native MOSS-Image support (PR #17095 at `76d935d`) using only stock nodes.
The DiT layers load as native quantized tensors (`int8_tensorwise` / `convrot_w4a4`).
On CPU, ComfyUI dequantizes instead of using the int8 / int4 GPU kernels. So this test confirms the files are valid and loadable,
not GPU kernel behavior, output quality or speed.
## ComfyUI usage
Two loaders work with these files. Use whichever matches your ComfyUI.
### A. Native ComfyUI core (no custom node)
Native MOSS-Image support is being added to ComfyUI core in [Comfy-Org/ComfyUI#17095](https://github.com/Comfy-Org/ComfyUI/pull/17095)
(open at the time of writing). With a ComfyUI build that includes it:
- **Load Diffusion Model** (`UNETLoader`): `moss-image-base-dit-int8-convrot.safetensors` or `...-int4-convrot.safetensors`, `weight_dtype = default`.
- **Load CLIP** (`CLIPLoader`): `moss-image-qwen35-int8-convrot.safetensors`, **type = `moss_image`**.
You can also use Comfy-Org's `qwen3.5_9b_int8_convrot.safetensors`, which has numerically equivalent weights.
- **Load VAE** (`VAELoader`): `moss-image-vae.safetensors`. Comfy-Org's `flux2-vae.safetensors` also works (see the comparison below).
- Workflows: `workflows/core/moss-image-base-int8-convrot.json` / `...-int4-convrot.json`, plus `*-api.json` for the HTTP API.
Core detects the separate q/k/v projections in these files automatically. Comfy-Org's own files use a fused `qkv`.
### B. Official MOSS-Image plugin
Requirements:
- A recent ComfyUI with comfy-kitchen >= 0.2.37. It must have INT4 ConvRot (`convrot_w4a4`) and ConvRot INT8 support; tested with ComfyUI `bdc63e7`.
- PyTorch built for CUDA 13.0 or newer, so the comfy-kitchen CUDA kernels are enabled.
- The official **MOSS-Image** plugin: search "MOSS-Image" in ComfyUI Manager, or copy
[`comfyui/moss-image`](https://github.com/OpenMOSS/MOSS-Image/tree/main/comfyui) into `ComfyUI/custom_nodes/`. Tested with plugin `1f808b7`.
On GPUs or devices without int8 compute, ComfyUI falls back to dequantize-then-matmul. This still saves memory but is not faster.
1. Download the files into your ComfyUI models directory, keeping the folder names:
```python
from huggingface_hub import hf_hub_download
for f in ["diffusion_models/moss-image-base-dit-int8-convrot.safetensors", # or ...-int4-convrot.safetensors
"text_encoders/moss-image-qwen35-int8-convrot.safetensors",
"vae/moss-image-vae.safetensors"]:
hf_hub_download("SumomoLee/MOSS-Image-Base-ConvRot-ComfyUI", f, local_dir="/path/to/ComfyUI/models")
```
2. Plugin workflows: load `workflows/moss-image-base-int8-convrot.json` or `workflows/moss-image-base-int4-convrot.json` in ComfyUI. The `*-api.json` files are for the HTTP API.
3. The nodes are **MOSS-Image Load Diffusion Model** (`MOSSImageUNETLoader`, keep `weight_dtype = default`), **MOSS-Image Load Qwen3.5**
(`MOSSImageCLIPLoader`), and the standard **Load VAE**.
Recommended settings for the base model, from the official user guide: 50 steps, CFG 3 (3 to 6; raise it for short prompts), shift 5 (ModelSamplingSD3),
Euler / sgm_uniform, uncond skip-layer guidance on layer 22 (SkipLayerGuidanceDiTSimple), and 1 to 2 MP resolutions in multiples of 16.
Recommended negative prompt: `nsfw, worst quality, bad anatomy, watermark, jpeg artifact`.
## Comparison with Comfy-Org/MOSS-Image
[Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image) ships BF16 and INT8 ConvRot DiTs for Base and SFT, an INT8 ConvRot Qwen3.5-9B
text encoder and `flux2-vae`. It does not ship INT4. Checked by reading tensors from both repos:
- **INT8 DiT:** the same `int8_tensorwise` + ConvRot (group 256, per-channel FP32 scale) format.
- Comfy-Org fuses q/k/v into `proj.qkv`; these files keep them separate, which both core and the plugin accept.
- The quantized values are identical up to a handful of single-step rounding ties (fewer than 1 in 5 million values).
- Per-layer weight error is identical (for example 0.855% vs 0.855% on qkv).
- Kept-precision tensors are byte-identical, including FP32 `rope.bands`.
- **Attention:** Comfy-Org's `*_int8_convrot_attn` files also carry `layers.N.comfy_attention.config = {"attention": "comfy_kitchen_int8"}`.
This opts each block into comfy-kitchen INT8 attention kernels where available. These files don't carry it, so attention runs at full precision:
slightly more accurate, possibly slower on GPUs with those kernels.
- **Text encoder:** the same quantization and layer selection.
- Comfy-Org's file is the stock Qwen3.5-9B checkpoint. It adds `lm_head` and `mtp.*`, which conditioning doesn't use, and stores `A_log` and `linear_attn.norm` in FP32.
- This file is quantized from MOSS-Image's own text encoder, which stores those two in BF16 (the FP32 values rounded).
- The quantized weights agree to within rare one-step rounding ties.
- Core prints a harmless "Missing weight for layer model.lm_head" warning for this file; MOSS-Image never uses `lm_head`.
- **VAE:** `moss-image-vae.safetensors` is the VAE shipped in MOSS-Image-Base. It is a FLUX.2 VAE with a slimmer decoder (96/192/384/384 channels, 62M params)
and slightly different encoder / batch-norm statistics from the stock `flux2-vae` (512-channel decoder, 84M params).
- The latent spaces are practically the same: under 0.4% difference in BN statistics and about 0.4% relative difference in encoded latents.
- The stock decoder reconstructs about 0.2–0.5 dB PSNR better in a quick test; the official decoder is smaller and faster.
- Either works.
## Changelog
- **v2:** the text encoder is re-keyed to `model.language_model.*` / `model.visual.*`. Tensor data is unchanged.
- The v1 key names (`language_model.*`, the plugin `export.py` naming) only load through the MOSS-Image plugin.
- In native ComfyUI core, v1 was silently misdetected as a CLIP-L model and produced invalid conditioning.
- v2 loads in both. Also: added core workflows, corrected the VAE description (FP32, not BF16), and added the Comfy-Org comparison. The DiT files are unchanged.
- **v1:** initial release.
## Checksums (SHA-256)
```text
1f42fac88a1b94f23ea3b49db6149243493ff1d6b962bb599a20417fe4aa71ea diffusion_models/moss-image-base-dit-int4-convrot.safetensors
f47c21a1397f846fa0d01c75e4b88fa6088783b47171399a465d92b1769ba0eb diffusion_models/moss-image-base-dit-int8-convrot.safetensors
511379fd690d59f11f49c062a8a687b9a985e92faa48a098eb7c66ba505574fe text_encoders/moss-image-qwen35-int8-convrot.safetensors
c3bdf94dc67ec69b30ea9fe1fa8340085679cae64d019f7c7e4a9921f1a7f17a vae/moss-image-vae.safetensors
```
## License and attribution
- **Model:** MOSS-Image-Base by the OpenMOSS Team ([Hugging Face](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base),
[GitHub](https://github.com/OpenMOSS/MOSS-Image)), revision `5f40f07a304541840b6019d2d34d30f966756787`, released under the Apache License 2.0.
These files are a modified (quantized) derivative of that model. A copy of the license is in [`LICENSE`](LICENSE). All credit for the model goes to the OpenMOSS Team.
- **Quantization recipe:** the INT4 ConvRot layout and layer selection follow
[toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI](https://huggingface.co/toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI).
The INT8 ConvRot DiT and text-encoder recipe and the workflow layout follow [Abiray/MOSS-Image-SFT-INT8-convrot](https://huggingface.co/Abiray/MOSS-Image-SFT-INT8-convrot).
Both use the ConvRot / comfy-kitchen layouts by Comfy Org.
- This is an unofficial community conversion, not affiliated with or endorsed by OpenMOSS. Check output quality on your own prompts before relying on it.