moss_pht

A lightweight, photo-oriented full fine-tune of 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

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

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 and user guide.

ComfyUI

Use ComfyUI's native MOSS-Image support with the 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: 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, copied from OpenMOSS-Team/MOSS-Image-Base).

We thank the OpenMOSS team for MOSS-Image.