--- base_model: OpenMOSS-Team/MOSS-Image-Base base_model_relation: adapter language: - en - zh license: apache-2.0 library_name: peft pipeline_tag: text-to-image tags: - moss-image - diffusers - lora - text-to-image - distillation --- # moss_image_base_dmd2 An 8-step LoRA (rank 128) for [MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/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](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. ![A goose resting by a rocky pond, generated in 8 steps](images/example_goose.png) *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](https://pytorch.org/get-started/locally/)), then: ```bash pip install "diffusers>=0.38" "transformers>=5.3" "accelerate>=1.10" "safetensors>=0.6" "peft>=0.21" "Pillow>=10" ``` ```python 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) and [user guide](https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md). ## License The weights 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.