Does Janus' image generation mode support multi-modal input?

[Dswn888]

Hi @learningpro @gnobitab @TheOneTrueGuy @mowentian @soloice , thanks for your excellent work!

I recently tried to use Janus for image generation.

In addition to providing text prompts, I also provided other images as part of the prompt. However, in the official implementation of Janus by transformers, I found the following code in the definition of JanusProcessor, which seems to mean that when generating images, the processor does not process the image elements in the input? So it seems that in image generation mode, the model only processes unimodal text prompts?

    # Process images if pixel values are provided.
    if images is not None and generation_mode != "image":
        data["pixel_values"] = self.image_processor(images=images, **output_kwargs["images_kwargs"])[
            "pixel_values"
        ]

Therefore, I hope to confirm with the official whether the image generation part of Janus supports multimodal input?

[avoroshilov]

I would like to understand this use case as well.
The Janus-Pro paper states in Figure 3 that, while for "Understanding" (effectively, text generation) and for "Image Generation", different image encoders should be used, the Image Generation conditioned on another image should be possible. The reason why transformers implementation has this limitation is exactly because you shouldn't pass an image through the "Understanding Encoder" for the image generation task.

Which means you need to pass it through the "Generation Encoder". Alas, there is not much detail on how exactly we do that.

Unfortunately, trying to do this manually doesn't yield any meaningful result. This script below works for regular Text2Image if I switch the DBG_SKIP_IMAGES = True which shows that it's generally working, but no matter how I construct the prompt, I cannot get meaningful images generated when injecting the image_embeds into the concatenated embeddings tensor.

# Prepare Prompts and Image
question = "Leave the image exactly as it is, but change color of the ladybug to blue."
image_path = "test_images/ladybug.png"
pil_image = PIL.Image.open(image_path).convert("RGB")

# Preprocess Image and get Visual Tokens (`Gen. Encoder` path)
image_transform = transforms.Compose([
	transforms.Resize(384),
	transforms.CenterCrop(384),
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
image_tensor = image_transform(pil_image).unsqueeze(0).to(vl_gpt.device, dtype=torch.bfloat16)
with torch.no_grad():
	_, _, (_, _, image_token_ids) = vl_gpt.gen_vision_model.encode(image_tensor)
image_tokens = image_token_ids.flatten() # Now a 1D tensor of visual tokens

# Define and Pre-embed all Prompt Building Blocks
# This creates a flexible "palette" of embedding tensors we can assemble.
tokenizer = vl_chat_processor.tokenizer
llm_embedder = vl_gpt.language_model.get_input_embeddings()

def get_embeds(text):
	tokens = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt").cuda().squeeze(0)
	return llm_embedder(tokens)

# Structural Components
bos_embed = llm_embedder(torch.LongTensor([tokenizer.bos_token_id]).cuda())
user_role_embed = get_embeds("<|User|>: ")
assistant_role_embed = get_embeds("<|Assistant|>: ")
separator_embed = get_embeds("\n\n")
boi_embed = get_embeds(vl_chat_processor.image_start_tag)
eoi_embed = get_embeds(vl_chat_processor.image_end_tag)

question_embed = get_embeds(question)
image_embeds = vl_gpt.prepare_gen_img_embeds(image_tokens)

DBG_SKIP_IMAGES = False

# Conditional Prompt Assembly
if DBG_SKIP_IMAGES:
	cond_parts = [
		bos_embed, user_role_embed,
		question_embed,
		#boi_embed, image_embeds, eoi_embed,
		separator_embed,
		assistant_role_embed, boi_embed
	]
else:
	cond_parts = [
		bos_embed, user_role_embed,
		boi_embed, image_embeds, eoi_embed,
		question_embed,
		separator_embed,
		assistant_role_embed, boi_embed
	]
cond_full_embeds = torch.cat(cond_parts, dim=0)

# Unconditional Prompt Assembly (no text question)
if DBG_SKIP_IMAGES:
	uncond_parts = [
		bos_embed, user_role_embed,
		# No question_embed
		#boi_embed, image_embeds, eoi_embed,
		separator_embed,
		assistant_role_embed, boi_embed
	]
else:
	uncond_parts = [
		bos_embed, user_role_embed,
		boi_embed, image_embeds, eoi_embed,
		# No question_embed
		separator_embed,
		assistant_role_embed, boi_embed
	]
uncond_full_embeds = torch.cat(uncond_parts, dim=0)

# Pad and Batch Embeddings
max_len = max(len(cond_full_embeds), len(uncond_full_embeds))
pad_embed = llm_embedder(torch.LongTensor([vl_chat_processor.pad_id]).cuda()).squeeze(0)

padded_cond = torch.cat([pad_embed.unsqueeze(0).repeat(max_len - len(cond_full_embeds), 1), cond_full_embeds], dim=0)
padded_uncond = torch.cat([pad_embed.unsqueeze(0).repeat(max_len - len(uncond_full_embeds), 1), uncond_full_embeds], dim=0)

initial_embeds = torch.stack([padded_cond, padded_uncond], dim=0)


### Do the generation loop similar to one in the "Text-to-Image Generation" `generate` function