| --- |
| 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` | |
| | --- | --- | |
| |  |  | |
|
|
| 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. |
|
|
|  |
|
|
| ## 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. |
|
|
