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

A lightweight, photo-oriented full fine-tune of
[OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base).
The text encoder (Qwen3.5-9B), VAE
(FLUX.2 small decoder), tokenizer, scheduler and inference pipeline are the base
release's, unchanged. This repository therefore ships only the transformer, in
two formats; load the other components from their official repositories, or
swap in any compatible ones (for example a quantized text encoder).

| File | Format |
| --- | --- |
| `transformer/` | Diffusers, the same layout as MOSS-Image-Base's `transformer/` |
| `diffusion_models/moss_pht_bf16.safetensors` | ComfyUI single file, the same layout as Comfy-Org's `moss_image_base_bf16.safetensors` |

## Diffusers

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

```python
import torch
from diffusers import AutoencoderKLFlux2
from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
from transformers import AutoModel, AutoTokenizer

BASE = "OpenMOSS-Team/MOSS-Image-Base"


def moss_class(name):
    return get_class_from_dynamic_module(
        BASE, "pipeline_moss_image.py", class_name=name, trust_remote_code=True
    )


pipe = moss_class("MOSSImagePipeline")(
    transformer=moss_class("MOSSImageTransformer").from_pretrained(
        "data-archetype/moss_pht", subfolder="transformer", torch_dtype=torch.bfloat16
    ),
    scheduler=moss_class("MOSSImageScheduler").from_pretrained(BASE, subfolder="scheduler"),
    text_encoder=AutoModel.from_pretrained(BASE, subfolder="text_encoder", dtype=torch.bfloat16),
    tokenizer=AutoTokenizer.from_pretrained(BASE, subfolder="tokenizer"),
    vae=AutoencoderKLFlux2.from_pretrained(BASE, subfolder="vae", torch_dtype=torch.float32),
)
pipe.to("cuda")  # or pipe.enable_model_cpu_offload() on a smaller GPU

image = pipe(
    "A candid street photograph of a cyclist crossing a rain-soaked intersection at dusk",
    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]
```

Each component is passed explicitly, so any of them can be replaced, for
example `text_encoder=` with a quantized Qwen3.5-9B. Sampling settings, prompt
formats (natural language, tags and JSON layouts) and requirements are the base
model's; see its [model card](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base)
and [user guide](https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md).

## ComfyUI

Use ComfyUI's native MOSS-Image support with the
[Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image) text encoder
and VAE; `moss_pht_bf16.safetensors` drops in for `moss_image_base_bf16.safetensors`
in the Comfy-Org workflows.

| File | Put in | Loader |
| --- | --- | --- |
| `diffusion_models/moss_pht_bf16.safetensors` (this repo) | `models/diffusion_models/` | Load Diffusion Model |
| `text_encoders/qwen3.5_9b_bf16.safetensors` (or its int8 / w4a8 variants), from Comfy-Org/MOSS-Image | `models/text_encoders/` | Load CLIP, type `moss_image` |
| `vae/flux2-vae_small.safetensors`, from Comfy-Org/MOSS-Image | `models/vae/` | Load VAE |

## Recommended: a DMD LoRA

Like the base checkpoint, generations with the standard multi-step sampler can be
a bit unstable. Using a DMD LoRA distilled from MOSS-Image-Base is recommended: it
sharpens the output distribution and stabilises generations.

## Training

| Setting | Value |
| --- | --- |
| Initialization | MOSS-Image-Base |
| Trained parameters | All transformer parameters (fp32 master weights, bf16 autocast); text encoder and VAE frozen |
| Steps | 12,500 at 64 images per step |
| Resolution | About 1 megapixel, aspect-ratio buckets |
| Text | Prompts up to 2,304 tokens; 10% of captions dropped to the empty prompt |
| Objective | The base's flow-matching objective; noise levels from a truncated logit-normal (μ 1.09, σ 0.95, truncated at 0.95), plus 1% uniform and 1% pure noise |
| Optimizer | [dionw](https://github.com/JTriggerFish/dionw): Dion (row-selected NorMuon) on the attention and MLP weight matrices (image blocks and the text embedder's blocks); AdamW on the input embeddings (latent, text and timestep), the AdaLN modulation projections, the output head, norms and biases |
| Learning rate | 9e-6, AdamW-equivalent for every parameter |
| Schedule | 200-step linear warmup, constant to step 10,500, then a 2,000-step 1 − √ cooldown to 0 |
| Dion matrices | 25% of each block's rows updated per step (blocks with both sides ≥ 1024), Gram Newton–Schulz, heavy-ball momentum 0.95, NorMuon row normalization β₂ 0.975 |
| AdamW parameters | β = (0.9, 0.975), ε 1e-8 |
| Weight decay | 0.01 on Dion matrices, none on AdamW parameters |
| Gradient clipping | Global norm 1.0 |
| Released weights | EMA of the weights (decay 0.999, fp32, updated every step) at step 12,500, stored in bfloat16 (`rope.bands` in float32, as in the base release) |

## License

The weights and the bundled inference code (`transformer/pipeline_moss_image.py`,
from OpenMOSS-Team/MOSS-Image-Base) are released under the base model's license,
Apache 2.0 ([LICENSE](LICENSE), copied from OpenMOSS-Team/MOSS-Image-Base).

We thank the OpenMOSS team for MOSS-Image.