thearkorbit/moss-image-edit-lora-poc

MOSS-Image Edit LoRA (proof of concept)

I got tired of waiting for an official editing version of MOSS-Image, so I opened up the model to see if it could edit images already.

Short version: MOSS-Image-SFT has a position-encoding axis that is never used. If you put a reference image there, the model can already see it with no training at all. This repo has a small proof-of-concept LoRA that starts teaching it to follow edit instructions (1 of 6 test edits worked so far).

Prior art: ostris/ai-toolkit already supports MOSS-Image edit LoRAs (extensions_built_in/diffusion_models/moss_image, added mid-October 2026). It puts reference latents on the same unused RoPE axis, and it also feeds the reference images to the Qwen-VL vision tower (krea2-style) and isolates references with a KV cache. I found the same RoPE trick on my own while poking at the model, before I knew about it. Their version is more complete, so use ai-toolkit for serious training. This repo is a small standalone write-up and experiment.

What I found

MOSS-Image-SFT is a 5B single-stream DiT. Its RoPE uses 4 coordinates per token, (r, f, h, w), but image tokens always get r = 0 (_RoPE.forward in transformer/pipeline_moss_image.py). The docstring also mentions "reference-image variants" that weren't released.

So the patch here:

1. Encodes the source image with the model's own FLUX.2 VAE (same 2x2 patchify + BN normalization the pipeline uses to decode).
2. Appends those latents as extra tokens after the noisy image, with RoPE r = 1.
3. Keeps output and loss on the noisy image tokens only.

With no adapter (ref_latents=None), the patched forward matches the original exactly (max abs diff 0.0).

Zero-shot (no training)

Test	Pixel L1 to source (lower = closer)
Same prompt, no reference	73.8
Same prompt, with reference	9.5
"Make the mug blue", with reference	9.4

The model copies the reference almost perfectly but ignores the edit instruction. So keeping the original image works out of the box. The LoRA only has to teach the editing part.

LoRA

• Rank 32 on q/k/v/o and FFN of all 42 blocks (74M trainable params)
• Data: 3,000 pairs from MagicBrush, 512px
• 1,200 steps, batch 4, AdamW 1e-4, logit-normal timesteps, 10% text dropout
• One A100 on Colab

Loss per 100 steps: 0.455 → 0.289 → 0.260 → 0.260 → 0.223 → 0.237 → 0.230 → 0.217 → 0.210 → 0.204 → 0.199 → 0.200

Results (honest version)

16 held-out MagicBrush pairs (not used in training), 28 Euler steps, CFG 3, one seed. These are the first 6, not cherry-picked:

• 1 of 6 edits landed (the fish on the surfboard). The other 5 still look like copies of the source.
• The LoRA fixes the colour shift the base model has with a reference (pinkish, washed out), so outputs stay much closer to the source.

Mean pixel L1 (0-255, lower = closer):

	to target	to source
base + ref (no LoRA)	20.2	16.9
LoRA	10.9	7.2
just copying the source	5.5	0

Pixel L1 can't tell editing from copying here: MagicBrush edits are small, so copying the source scores best. The pictures are the better evidence.

Takeaway: the reference path works and the model can pick up edits from very little data, but 3K pairs and ~1.6 epochs isn't enough. The copying habit still wins most of the time.

Usage

import torch
from diffusers import DiffusionPipeline
from moss_edit import patch, encode_image, load_lora

pipe = DiffusionPipeline.from_pretrained(
    "OpenMOSS-Team/MOSS-Image-SFT", custom_pipeline="pipeline_moss_image",
    trust_remote_code=True, torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
).to("cuda")
patch(pipe.transformer)
load_lora(pipe.transformer, "moss_edit_lora_poc.safetensors")
ref = encode_image(pipe.vae, source_image, (512, 512))
pipe.transformer._ref = ref.to(torch.bfloat16)
image = pipe("make the sky purple", height=512, width=512).images[0]

Limitations

This is a proof of concept, not a finished editor. It was trained on 3K pairs for about an hour, at 512px only (MOSS-Image is trained at 1024+). Most edits don't land yet.

Next: more and cleaner edit data, longer training, 1024px, and reference dropout so the model stops leaning on copying.