Nodes/ZenTrainer/MOSS Edit Model (references)
ComfyUI Node

MOSS Edit Model (references)

One wire, one input, and MOSS Image finally reads your photo

By Nynxz·Created 8 days ago·Updated about 23 hours ago· 5
MOSS Edit Model (references)
  • model
  • MODEL

One input, one output, zero settings. It looks like the laziest node in the pack and it's the load-bearing one: without it, reference-picture edits on MOSS Image don't happen at all.

Why a node this small has to exist

Training an edit LoRA in ai-toolkit means the references join the sequence as clean tokens - the "reference latents" arrangement Kontext and Krea 2 use. When you sample, the model has to see them the same way or the LoRA is being asked to perform a trick it never learned.

Here's the problem: ComfyUI's MOSS Image implementation doesn't read reference_latents from conditioning - upstream support exists as a pull request (ComfyUI #17296), not shipped behaviour. So the pack patches it at the model level, which is exactly what a MODEL-in, MODEL-out node is for. That plumbing habit of "wrap the patcher, hand it on" is why this is a node and not a toggle on the encode node.

What it changes on the patcher

Reading harness/moss_refs.py rather than guessing: it clones your patcher and then does three things.

First, it teaches extra_conds to carry reference_latents through as ref_latents, plus the chosen method. That's the handoff - the conditioning from MOSS Edit Encode (or from ComfyUI's own ReferenceLatent node) becomes something the diffusion model actually receives.

Second, it replaces the diffusion forward with MOSS's own sequence, plus the reference tokens on their own rope axis. Each reference sits on its own value of the rope axis MOSS leaves at zero (r = 1, 2, …), at the picture's frame position, rows and columns centered like the picture's. The references are therefore aligned with your output geometry but distinguishable from it. That's the whole trick, and it's the same layout the upstream PR uses, so a LoRA trained in ai-toolkit samples the same way here.

Third, the two layouts. In index_timestep_zero the references get their own pass at timestep zero, each block's keys and values are kept, and then the picture's pass attends to them alongside its own - the ai-toolkit kv-cache arrangement. In index the reference tokens ride after the text at the current step's noise level and attend to everything. Whichever your LoRA was trained with, this node makes it available at sampling time.

Worth knowing: the reference pass deliberately skips ComfyUI's block patches, so skip-layer-guidance applies only to the picture's pass. That's correct - SLG belongs to your prompt batch, not the reference tokens.

Putting it in the graph

model in, MODEL out. Put it after your LoRA loader, before the sampler - the author's instruction, and it means the patch lands on the final patcher with your LoRA already baked in.

The line looks like: checkpoint or UNET loader → LoRA loader → MOSS Edit Model (references) → sampler. Conditioning comes from MOSS Edit Encode, whose positive/negative/latent feed the guider and sampler. Skipping the encode node because you already have text conditioning? ComfyUI's ReferenceLatent node sets the same conditioning key and also works - you'll get index_timestep_zero by default.

Sampling settings that match the model's own guide: Euler, 50 steps, CFG 3, layer 22 skipped in the negative pass.

Install

Manager → search ZenTrainer, or:

cd ComfyUI/custom_nodes
git clone https://github.com/Nynxz/ZenTrainer
# restart

Then look under zentrainer/moss in the node search. ComfyUI 0.41 or newer is required and the pack says so if yours is older. No extra pip installs for this node - requirements.txt notes everything it needs ships with ComfyUI. You will need MOSS Image itself (diffusion model, Qwen3.5 text encoder, VAE) downloaded, which this node can't do for you.

Where it goes wrong

Feeding it a non-MOSS model. It checks, and it's a clear error rather than a crash: MOSS references are for MOSS Image, not <something else>. You can't retrofit reference support onto Flux or SDXL with this - the rope layout is MOSS-specific.

Expecting it to work without references in the conditioning. The node itself is inert if nothing in your graph supplies reference latents. Picture it as a socket, not a source.

Assuming it's permanent. This is a shim over an upstream pull request. Once ComfyUI reads reference latents for MOSS natively the node is probably redundant - but until then it's the only way to sample an edit LoRA trained with references instead of channel-input pairs.

Speed and VRAM. The default layout runs the reference pictures through the blocks once, on their own, before the picture does. Two or three references isn't free; if a generation got noticeably slower after you added image2 and image3, that's the reason, not a bug.

Mixing it with the wrong LoRA. A LoRA trained on channel-input edit pairs (the other route ZenTrainer supports) won't thank you for reference latents. If edits come out as a loose restyle of the input rather than the change you asked for, check how the LoRA was trained before touching the sampler.

Categoryzentrainer/moss

Inputs (1)

NameTypeDefaultDescription
modelMODEL—

Outputs (1)

NameTypeDescription
MODELMODEL—