moss_image_base_dmd2

An 8-step LoRA (rank 128) for MOSS-Image-Base, distilled by distribution matching. Inference is 8 forwards at CFG 1, without a negative prompt, in place of the base model's 50 guided steps. Training details 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.

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

Sampling

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

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.

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.