Keeping the same face across generations: how reference photos enter an OmniGen prompt

You want to generate a new picture of a specific person. Not a generic face, that exact person, the one in your three reference photos. You write a prompt like "the woman in <|image_1|> standing in a kitchen" and you expect the face to survive. Half the time it does not. You get someone who is sort of in the neighborhood, a cousin maybe.

The interesting part of identity preservation in a model like OmniGen is not the diffusion. It is the few dozen lines that decide where your reference photos go inside the prompt and how the model is told "these pixels are the person, the rest is instructions." That happens in the processor, before a single denoising step runs. This is a walkthrough of that code path in the Wunjo Make repo.

OmniGen (Xiao et al., 2024) is a single transformer that takes text and images mixed together in one sequence, instead of a separate text encoder feeding a separate U-Net. That design is exactly why the conditioning code is worth reading: text tokens and image features live in the same stream, and the splice between them is where identity is kept or lost.

Several reference photos and one prompt all become one flat sequence of tokens. Photo: Unsplash.

The prompt is a string with image holes in it

You reference images by tag. The prompt literally contains <|image_1|>, <|image_2|>, and so on. The first job is to split the text at those tags and tokenize the text between them. From processor.py:

pattern = r"<\|image_\d+\|>"
prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)]

for i in range(1, len(prompt_chunks)):
    if prompt_chunks[i][0] == 1:
        prompt_chunks[i] = prompt_chunks[i][1:]

re.split on the tag pattern gives you the text pieces in order, with the tags removed. Each piece is tokenized on its own. That second loop is a small but real detail: the tokenizer prepends a beginning-of-sequence token (id 1) to every chunk, and you only want one at the very start of the whole prompt. So every chunk after the first gets its leading 1 stripped. Miss this and you sprinkle BOS tokens through the middle of your sequence, which is not what the model was trained on.

Tags carry an order, and the order is enforced

Next the code reads the tags back out and turns them into integer ids:

image_tags = re.findall(pattern, text)
image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]

unique_image_ids = sorted(list(set(image_ids)))
assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), \
    f"image_ids must start from 1, and must be continuous int, ... cannot be {unique_image_ids}"
assert len(unique_image_ids) == len(input_images), \
    f"total images must be the same as the number of image tags, ..."

input_images = [input_images[x-1] for x in image_ids]

Two assertions do real work here. The ids must be [1, 2, 3, ...] with no gaps, and the number of distinct tags must equal the number of images you passed. If you write <|image_1|> and <|image_3|> but skip 2, it stops. If you pass two photos but only tagged one, it stops. This is the kind of guard you want, because the failure it prevents is silent: the model would happily generate from a misaligned set and you would spend an hour wondering why the face is wrong.

That last line is the one to sit with. input_images = [input_images[x-1] for x in image_ids] reorders your actual image tensors to match the order the tags appear in the text. The position in the sentence, not the position in your list, is what counts. A tag can repeat too, and the same photo gets placed at each spot it is named.

Reserving the right number of slots for each image

Now the sequence gets assembled. Text token ids and image placeholders are concatenated into one flat list:

all_input_ids = []
img_inx = []
idx = 0
for i in range(len(prompt_chunks)):
    all_input_ids.extend(prompt_chunks[i])
    if i != len(prompt_chunks) - 1:
        start_inx = len(all_input_ids)
        size = input_images[i].size(-2) * input_images[i].size(-1) // 16 // 16
        img_inx.append([start_inx, start_inx + size])
        all_input_ids.extend([0] * size)

return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx}

Walk it. After each text chunk (except the last) there is an image, so the code computes how many tokens that image will occupy and pushes that many placeholder 0s into the sequence. The count is H * W // 16 // 16. An image is encoded by the VAE down by a factor of 8, then split into patches of size 2 by the patch embedder, so each side shrinks by 16 total and one image becomes (H/16) * (W/16) tokens. A 1024 by 1024 photo is 64 * 64 = 4096 tokens. That is a lot of sequence per image, which is also why three references cost real memory.

The 0s are pure reservation. They are placeholders, not the image. What matters is img_inx, the list of [start, end] spans recording exactly where each image sits in the flat sequence. Those spans are the address book the model uses later to drop the real image features into place.

The placeholder spans become image features

The actual splice happens in model.py, inside forward. The input ids go through the text embedding table as usual, then the recorded spans are overwritten with VAE-encoded, patch-embedded image latents:

condition_embeds = self.llm.embed_tokens(input_ids).clone()
input_img_inx = 0
for b_inx in input_image_sizes.keys():
    for start_inx, end_inx in input_image_sizes[b_inx]:
        condition_embeds[b_inx, start_inx:end_inx] = input_latents[input_img_inx]
        input_img_inx += 1

This is the whole trick. The [0] placeholders were embedded to some throwaway vectors; this loop replaces those vectors, span by span, with the features of the actual reference photos. After it runs, condition_embeds is one sequence where some positions are text and some positions are "here is exactly what this person's face looks like, at this patch." Then time embedding and the noisy latent are concatenated on, and the whole thing goes through the Phi-3 transformer backbone (Abdin et al., 2024) that OmniGen builds on. Text and image are processed by the same attention layers, in the same sequence. That shared attention is how a token of the prompt can look directly at a patch of the reference face.

One more thing: image tokens see each other fully

There is a detail in the collator that matters for identity. The base attention mask is causal (each token sees only earlier tokens), but for the image spans that is wrong; a patch in the top-left of a face should be able to see a patch in the bottom-right. So the collator opens up full attention inside each image block:

def adjust_attention_for_input_images(self, attention_mask, image_sizes):
    for b_inx in image_sizes.keys():
        for start_inx, end_inx in image_sizes[b_inx]:
            attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1
    return attention_mask

Inside a reference image, every patch attends to every other patch. Between images and text, the causal order still holds. So each photo is read as a coherent whole, not as a left-to-right scan, which is what you want when the thing you are trying to preserve is a face.

Gotchas you will actually hit

A few things that bite in practice, all visible in the code above:

Tag your images, do not just pass them. If your instruction does not contain <|image_1|>, the assertion path is skipped and the photos are never spliced in. Passing images to the pipeline is not enough; the tag is the anchor.
Number from 1, no gaps. [1, 2, 3], never [0, 1] or [1, 3]. The assertion is strict on purpose.
Reference resolution is sequence length. Each 1024 image is 4096 tokens before your prompt and the output canvas are even counted. If you are running out of memory with three references, that is why. Smaller reference crops cost fewer tokens.
Order is semantic. Because images are reordered to tag order, "<|image_1|> wearing the hat from <|image_2|>" is not the same prompt as swapping the two tags. Be deliberate about which photo is which number.

About the author. I'm Wlad Radchenko, a software engineer. The code in this article comes from Wunjo Make (open source), local software for video makers, and Wunjo Design, an offline PWA for designers. Get in touch to find more on GitHub and LinkedIn.

Takeaway

Identity preservation across generations is not magic in the sampler. It is bookkeeping in the processor: split the prompt at image tags, tokenize the text between, reserve H*W//16//16 placeholder slots per image, record the span of each, then overwrite those spans with real image features so text and face share one attention stream. The model can only keep the same person if the same person's pixels are sitting in the sequence, in known positions, with full attention inside each photo.

The code is in visual_generation/generation/omnigen/processor.py and model.py in the Wunjo Make repo. If your reference-based results drift, start by printing image_sizes and confirming the spans line up with the images you think you passed.