moss_image_base_dmd2

An 8-step LoRA (rank 128) for MOSS-Image-Base, distilled by distribution matching: 8 forwards at CFG 1, without a negative prompt, in place of the base model's 50 guided steps. It adapts only the transformer; the text encoder (Qwen3.5-9B), VAE (FLUX.2 small decoder), tokenizer and inference pipeline are MOSS-Image-Base's, unchanged. How it was trained is in TECH_REPORT.md.

It was trained on MOSS-Image-Base, but should work with most fine-tunes of it and with other LoRAs: load it alongside them, at strength 1.

8 steps at CFG 1, 1152 × 896, seed 42, with the prompt in the Diffusers example below. The image carries its ComfyUI workflow: drag it into ComfyUI to load it.

File	Format
loras/moss_image_base_dmd2_bf16.safetensors	ComfyUI LoRA (kohya keys on ComfyUI's module names, q/k/v as one LoRA on the fused qkv)
transformer/	PEFT LoRA adapter for the Diffusers MOSSImageTransformer
images/example_goose.png	Example image, its ComfyUI workflow embedded
workflows/moss_image_base_dmd2_8step.json	ComfyUI workflow (drag into ComfyUI)
workflows/moss_image_base_dmd2_8step_api.json	The same workflow in ComfyUI's API format

Both LoRA files make the same weight change; the ComfyUI file packs each block's q, k and v LoRAs into one exactly equivalent LoRA on ComfyUI's fused qkv projection.

Sampling

The LoRA is trained for one sampler; others give worse results.

Setting	Value
Steps	8
Sampler	At each step, take the model's clean estimate, then re-noise it with fresh noise to the next level (ComfyUI's lcm)
Schedule	8 evenly spaced times 1, 7/8, ..., 1/8, then 0, under the flow shift 4.4817 (ComfyUI's simple scheduler after a ModelSamplingAuraFlow node at shift 4.4817); about 1, .969, .931, .882, .818, .729, .599, .39
CFG	1, no negative prompt (guidance is distilled in)
LoRA strength	1
Resolution	About 1 megapixel, any aspect ratio the base supports

ComfyUI

Use ComfyUI's native MOSS-Image support and load the workflow: drag images/example_goose.png or workflows/moss_image_base_dmd2_8step.json into ComfyUI.

File	Put in	Loader
loras/moss_image_base_dmd2_bf16.safetensors (this repo)	models/loras/	LoraLoaderBypassModelOnly
diffusion_models/moss_image_base_bf16.safetensors, from Comfy-Org/MOSS-Image	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

The workflow sets the sampler (lcm), scheduler (simple), 8 steps, CFG 1 and the shift. ComfyUI's default MOSS-Image shift is 5, close to but not the 4.4817 the LoRA was trained with; the workflow's ModelSamplingAuraFlow node sets it. Use that node (or ModelSamplingFlux) to change the shift, not ModelSamplingSD3: it scales MOSS-Image's timesteps differently and the output turns to noise.

Load the LoRA with the bypass loader ("Load LoRA (Bypass, Model Only)"), as the workflow does, not the standard LoRA loader. The standard loader merges the LoRA into the bf16 weights, and the weight change is smaller than bf16's rounding step, so most of it is lost. The bypass loader applies it unmerged, as Diffusers does below.

Diffusers

Install a PyTorch build that matches your CUDA driver first (pytorch.org), then:

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

import math

import torch
from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler
from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
from peft import PeftModel
from transformers import AutoModel, AutoTokenizer

BASE = "OpenMOSS-Team/MOSS-Image-Base"
SHIFT = math.exp(1.5)


class FreshNoiseScheduler(FlowMatchEulerDiscreteScheduler):
    """The LoRA's sampler: jump to the clean estimate, re-noise with fresh noise."""

    def set_timesteps(self, num_inference_steps=8, device=None):
        t = torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
        sigmas = SHIFT * t / (1 + (SHIFT - 1) * t)
        super().set_timesteps(device=device, sigmas=sigmas.tolist())

    def step(self, model_output, timestep, sample, return_dict=True, **kwargs):
        if self.step_index is None:
            self._init_step_index(timestep)
        sigma, sigma_next = self.sigmas[self.step_index], self.sigmas[self.step_index + 1]
        clean = sample.float() - sigma * model_output.float()
        if sigma_next > 0:
            clean = (1 - sigma_next) * clean + sigma_next * torch.randn_like(clean)
        self._step_index += 1
        return (clean.to(sample.dtype),)


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


transformer = moss_class("MOSSImageTransformer").from_pretrained(
    "OpenMOSS-Team/MOSS-Image-Base", subfolder="transformer", torch_dtype=torch.bfloat16
)
transformer = PeftModel.from_pretrained(
    transformer, "data-archetype/moss_image_base_dmd2", subfolder="transformer"
).get_base_model()

pipe = moss_class("MOSSImagePipeline")(
    transformer=transformer,
    scheduler=FreshNoiseScheduler(shift=1.0),
    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

torch.manual_seed(42)
image = pipe(
    "A large brown and white goose resting on dry grass at the edge of a rocky pond, its patterned wing feathers folded neatly and its long neck curved as it looks toward the water. The sun casts warm light across the bird and the rough stone bank behind it.",
    height=896,
    width=1152,
    num_inference_steps=8,
    guidance_scale=1.0,
    uncond_skip_layers=None,
).images[0]

The adapter is applied unmerged, at strength 1. Prompt formats (natural language, tags and JSON layouts) and requirements are the base model's; see MOSS-Image-Base's model card and user guide.

License

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