Nodes/Anzhc-ComfyUI-Cosmos-Reference/Apply Cosmos Reference Model Patch
ComfyUI Node

Apply Cosmos Reference Model Patch

The boring half of the pair — skip it and the reference silently does nothing

By Anzhc·Created 4 months ago·Updated 4 months ago· 8
Apply Cosmos Reference Model Patch
  • model
  • MODEL

This is the node nobody searches for and everybody in the Anima control workflow needs. Its sibling, Cosmos Reference Conditioning, does the easy part - it tags your conditioning with a reference latent and a strength. This node does the surgery: it patches the Cosmos model so it can actually read that metadata during sampling. Skip it, and your carefully-wired reference is silently ignored, which is exactly the "strength does nothing" report people hit.

How it works

One MODEL in, one patched MODEL out. Wire it between your UNETLoader (or LoraLoader, as the repo's Anima Control workflow does) and the KSampler. Under the hood it clones the model and applies four patches:

  1. An extra_conds override so the reference latents, strength and control method flow through the conditioning pipeline.
  2. A prepare_embedded_sequence patch that gives temporal-mask models their padding-mask channel.
  3. Per-block attention op swaps that enable KV gating on reference tokens.
  4. A model function wrapper that concatenates reference latents onto the temporal axis, aligns the sequence to the model's temporal patch size, then slices the output back to the original length - so the patch is transparent to everything else in the graph.

Because it's all applied to a cloned model, your base model stays untouched and the patch is isolated to this sampler run.

Gotchas

  • It does nothing alone. The patch node has no reference to consume unless CosmosReferenceConditioning stamped the conditioning. They're a pair; both need to be in the graph, and the patched model has to be the one the KSampler actually uses.
  • Temporal mask is picky about the model. If you pick temporal_mask on a model without concat_padding_mask=True, it raises a RuntimeError rather than degrading silently. The Anima / Cosmos-Predict2-2B line has it; other architectures don't.
  • KV gating has an off switch. At strength 0 the reference is dropped entirely - which the author points to as a feature, since it means a control LoRA can be present in the graph without forcing its reference.
  • Cosmos only. This is built on the Cosmos attention and embedding structure. It won't do anything useful on SDXL, Flux, or Illustrious - and honestly, those models don't need it; they have ControlNet.

Install

Same as the rest of the pack, and it's a clean one:

cd ComfyUI/custom_nodes
git clone https://github.com/Anzhc/Anzhc-ComfyUI-Cosmos-Reference

Restart ComfyUI. No requirements.txt, no extra Python deps - just torch, einops and ComfyUI internals, all already present. ComfyUI Manager users can search "Anzhc-ComfyUI-Cosmos-Reference" instead. Nothing to download from the pack itself; you supply the Cosmos model and, for real control behavior, the companion LoRA.

Where it sits

This is the infrastructure half of a two-node answer to Anima's missing ControlNet problem. Anzhc built it by taking Mirumo0u0's original Cosmos reference patch and extending it with proper strength control - the kind of thing you'd never install on its own, but without which the whole reference conditioning approach doesn't function. It's not flashy. It just needs to be there, in the right place, between your model and your sampler.

CategoryCosmos/Reference

Inputs (1)

NameTypeDefaultDescription
modelMODEL

Outputs (1)

NameTypeDescription
MODELMODEL