File size: 7,199 Bytes
969b5a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | ---
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.
*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.
|