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