File size: 6,807 Bytes

6c483fd

---
license: apache-2.0
base_model: OpenMOSS-Team/MOSS-Image-Base
base_model_relation: finetune
language:
- en
pipeline_tag: text-to-image
tags:
- moss-image
- diffusers
- text-to-image
---
# Kromoss

Kromoss is a full fine-tune of [MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) at
1024 px. It uses the same architecture, inference pipeline, text encoder (Qwen3.5-9B) and VAE (FLUX.2 small
decoder) as [MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT), and this repo has the same
layout, so it is a drop-in replacement. The fine-tune mainly fixes the Base model's blurriness.

The repo also ships `8-steps-dmd-lora.safetensors`, an 8-step turbo LoRA (data-space DMD distillation of
MOSS-Image-Base): **8 steps, CFG 1**, no negative branch.

| 50 steps, CFG 3 | 8 steps, CFG 1 + `8-steps-dmd-lora` |
| --- | --- |
| ![50 steps](assets/cat-50steps-cfg3.webp) | ![8 steps + LoRA](assets/cat-8steps-lora.webp) |

Checkpoint: update 5000 of the fine-tune (EMA weights), 320k training images.

## Quick start

```bash
pip install "diffusers>=0.38" "transformers>=5.3" "accelerate>=1.10" "safetensors>=0.6" "Pillow>=10"
```

```python
import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "lodestones/Kromoss",
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True,
    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
)
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
pipe.to("cuda")

image = pipe(
    "a tabby cat holding a sign that says KROMOSS",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=3.0,
    uncond_skip_layers=[22],
    generator=torch.Generator("cpu").manual_seed(42),
).images[0]
image.save("kromoss.png")
```

The transformer and text encoder use bfloat16; the VAE and initial noise use float32.

## ComfyUI

Native ComfyUI support for MOSS-Image is in current ComfyUI master. The single files use the same layout as
[Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image): fused `proj.qkv`, bf16.

| File | Put in | Loader |
| --- | --- | --- |
| `diffusion_models/kromoss_bf16.safetensors` | `models/diffusion_models/` | Load Diffusion Model |
| `diffusion_models/kromoss_int8_convrot.safetensors` | `models/diffusion_models/` | Load Diffusion Model (int8, 5.4 GB) |
| `text_encoders/qwen3.5_9b_bf16.safetensors` | `models/text_encoders/` | Load CLIP, type `moss_image` |
| `vae/flux2-vae_small.safetensors` | `models/vae/` | Load VAE |
| `8-steps-dmd-lora.safetensors` | `models/loras/` | LoraLoaderModelOnly, strength 1 |

`kromoss_bf16.safetensors` drops in for `moss_image_sft_bf16.safetensors` in the
Comfy-Org workflows (50 steps, CFG 3). With the LoRA, use KSampler with `euler` / `simple`, 8 steps, CFG 1.
ComfyUI maps the LoRA's separate q/k/v onto the fused `qkv` slices.

`convert_to_comfy.py` converts any MOSS-Image diffusers `transformer/` folder (e.g. your own fine-tune) into this
single-file format (needs only `torch` and `safetensors`):

```bash
python convert_to_comfy.py path/to/transformer my_model_bf16.safetensors
```

On MOSS-Image-SFT's `transformer/` it reproduces Comfy-Org's `moss_image_sft_bf16.safetensors` byte for byte;
`--compare FILE` checks keys, shapes and dtypes against a reference file.

`kromoss_int8_convrot.safetensors` is ConvRot W8A8, the format of Comfy-Org's `moss_image_sft_int8_convrot.safetensors`.
The 168 block linears (`proj.{qkv,o,ffn_up_gate,ffn_down}`) are stored as int8 with per-row scales, after a
256-wide regular Hadamard rotation of the input dimension. ComfyUI rebuilds the rotation at load time and applies it
to the activations online, so no rotation matrices are stored. Each block also selects ComfyUI's int8 attention.
It drops in for the int8 Comfy-Org file and works with the LoRA. `quantize_convrot.py` makes it from a bf16 single
file; on `moss_image_sft_bf16.safetensors` it reproduces Comfy-Org's int8 file byte for byte:

```bash
python quantize_convrot.py my_model_bf16.safetensors my_model_int8_convrot.safetensors
```

## 8-step turbo LoRA

`8-steps-dmd-lora.safetensors`: fp32, rank 128 / alpha 128 on every block's `proj.{q,k,v,o,ffn_up_gate,ffn_down}`,
kohya keys (`diffusion_model.<module>.lora_down.weight` / `lora_up.weight` / `alpha`). In ComfyUI, load it with
a LoRA loader at strength 1. With diffusers, apply it unmerged. Do not merge it into the bf16 weights: the delta
is close to bf16's rounding step.

```python
from safetensors.torch import load_file

def apply_lora(transformer, path, strength=1.0):
    state = load_file(path, device="cuda")
    modules = dict(transformer.named_modules())
    for key in [k for k in state if k.endswith(".lora_down.weight")]:
        name = key[len("diffusion_model."):-len(".lora_down.weight")]
        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()
        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]
        def hook(module, args, out, down=down, up=up, scale=scale):
            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)
        modules[name].register_forward_hook(hook)

apply_lora(pipe.transformer, "8-steps-dmd-lora.safetensors")   # hf_hub_download("lodestones/Kromoss", ...)
image = pipe(prompt, height=1024, width=1024, num_inference_steps=8, guidance_scale=1.0,
             generator=torch.Generator("cpu").manual_seed(42)).images[0]
```

## Sampling settings

| Setting | Full model | With `8-steps-dmd-lora` |
| --- | --- | --- |
| Sampler | Flow-matching Euler | Flow-matching Euler |
| Steps | 50 | 8 |
| CFG | 3 | 1 |
| Shift | 5 | 5 |
| Layer skipping | Negative branch, layer index 22 | none (no negative branch) |

## PartiPrompts (100 prompts, CLIP ViT-L/14 score; Kromoss rows measured at update 2000)

| Model | CLIP |
| --- | --- |
| MOSS-Image-Base, 50 steps, CFG 3 | 27.19 |
| MOSS-Image-Base + LoRA, 8 steps, CFG 1 | 27.53 |
| Kromoss, 50 steps, CFG 3 | 26.88 |
| Kromoss + LoRA, 8 steps, CFG 1 | **27.82** |

Columns in the sheet below follow the table's order.

![PartiPrompts comparison](assets/parti-comparison.webp)

## Training

- Full fine-tune of every transformer parameter, 1024 px (images under 896² dropped).
- The MOSS-Image report's final-stage recipe, scaled to 64 images per update:
  - hybrid Muon-AdamW, lr 9e-6, warmup 200;
  - truncated logit-normal timesteps;
  - final-block q/k-norm decay to one;
  - EMA 0.999 (the published weights are the EMA).

## License and acknowledgements

The weights and the bundled inference code (`pipeline_moss_image.py`, from MOSS-Image-SFT) are released under the
[Apache License, Version 2.0](LICENSE). Built on [MOSS-Image](https://github.com/OpenMOSS/MOSS-Image) by the
OpenMOSS Team.