Emu3: add model

[zucchini-nlp]

What does this PR do?

As per title. The code can work for generating text in single-batch scenarios but the generated text doesn't match input image. For batched generation, seems like the orig impl neither supports it mostly because image features from processor are returned with different shapes (smart resize to converse as much orig image size as possible). We can try to do padding similar to llava-next but I am not sure if will just work, I'll contact the authors

TODO:

• Batched generation
• Upload chat template and change the image-placeholder token from extra-0 to smth like <image>
• Match the orig implementation on logit level
• Tests, many more tests
• Check out image generation and see how we can enable interleaved image+text generation as in Chameleon. Maybe not natively with transformers, but we can provide scripts with external libraries for structured generation -> not possible because text-generation and image-generation are two different checkpoints with different weights

from PIL import Image
import torch
import requests

from transformers import (
    Emu3Config,
    Emu3ForConditionalGeneration,
    Emu3ImageProcessor,
    Emu3Processor,
)

output_dir = "/raid/raushan/emu3"
processor = Emu3Processor.from_pretrained(output_dir)
model = Emu3ForConditionalGeneration.from_pretrained(output_dir, torch_dtype="bfloat16", device_map="auto")
processor.tokenizer.padding_side = "left"

text = "You are a helpful assistant. USER: <|extra_0|>Please describe the image. ASSISTANT:"
image = Image.open("/raid/raushan/image.png")
image2 = Image.open(requests.get("https://www.ilankelman.org/stopsigns/australia.jpg", stream=True).raw)

inputs = processor(
    text=[text, text],
    images=[image2, image],
    return_tensors="pt",
    padding=True,
)

inputs = inputs.to(device="cuda:0", dtype=torch.bfloat16)

out = model.generate(**inputs, max_new_tokens=100)
text_out = processor.batch_decode(out, skip_special_tokens=True)
print(text_out)

And for image generation:

from PIL import Image
from transformers import AutoTokenizer, AutoModel, AutoImageProcessor, AutoModelForCausalLM
import torch
import requests

from transformers import (
    Emu3Config,
    Emu3ForConditionalGeneration,
    Emu3ImageProcessor,
    Emu3Processor,
)

output_dir = "/raid/raushan/emu3-gen"
processor = Emu3Processor.from_pretrained(output_dir)
model = Emu3ForConditionalGeneration.from_pretrained(output_dir, torch_dtype="bfloat16", device_map="auto", ) # attn_implementation="flash_attention_2",


inputs = processor(
    text=["a portrait of young girl. masterpiece, film grained, best quality.", "a dog running under the rain"],
    padding=True,
    return_tensors="pt",
    return_for_image_generation=True,
)
inputs = inputs.to(device="cuda:0", dtype=torch.bfloat16)

image_sizes = inputs.pop("image_sizes")
HEIGHT, WIDTH = image_sizes[0]
VISUAL_TOKENS = model.model.vocabulary_mapping.image_tokens

def prefix_allowed_tokens_fn(batch_id, input_ids):
    height, width = HEIGHT, WIDTH
    visual_tokens = VISUAL_TOKENS
    image_token_id = processor.tokenizer.encode("<|image token|>", return_tensors="pt")[0].to(model.device) # torch.tensor([processor.tokenizer.image_token_id], device=model.device)
    eoi_token_id = processor.tokenizer.encode("<|image end|>", return_tensors="pt")[0] # torch.tensor([processor.tokenizer.eoi_token_id], device=model.device)
    eos_token_id = processor.tokenizer.encode("<|extra_204|>", return_tensors="pt")[0] # torch.tensor([processor.tokenizer.eos_token_id], device=model.device)
    pad_token_id = processor.tokenizer.encode("<|endoftext|>", return_tensors="pt")[0] # torch.tensor([processor.tokenizer.pad_token_id], device=model.device)
    eol_token_id = processor.tokenizer.encode("<|extra_200|>", return_tensors="pt")[0]
    eof_token_id = processor.tokenizer.encode("<|extra_201|>", return_tensors="pt")[0]

    position = torch.nonzero(input_ids == image_token_id, as_tuple=True)[0][0]
    offset = input_ids.shape[0] - position
    if offset % (width + 1) == 0:
        return (eol_token_id, )
    elif offset == (width + 1) * height + 1:
        return (eof_token_id, )
    elif offset == (width + 1) * height + 2:
        return (eoi_token_id, )
    elif offset == (width + 1) * height + 3:
        return (eos_token_id, )
    elif offset > (width + 1) * height + 3:
        return (pad_token_id, )
    else:
        return visual_tokens


out = model.generate(
    **inputs,
    max_new_tokens=50_000,
    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
    do_sample=True,
    top_k=2048,
    return_dict_in_generate=True,
)

print(out.sequences.shape, inputs.input_ids.shape)

image = model.model.decode_image_tokens(out.sequences[:, inputs.input_ids.shape[1]: ], height=HEIGHT, width=WIDTH)
images = processor.postprocess(list(image.float()), return_tensors="PIL.Image.Image") # internally we convert to np but it's not supported in bf16 precision
for i, image in enumerate(images['pixel_values']):
    image.save(f"result_{i}.png")

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zucchini-nlp]

Adding LlamaModel to bases messes up the auto-generated modeling file by adding new classes like Emu3TextAttention and so on, while we have Emu3Attention

[Cyrilvallez]

I think this should be solved by #34487!

[ArthurZucker]

To try again!

[zucchini-nlp]

I think this is ready for review. @ArthurZucker will you be reviewing or is there anyone I can tag for initial review?

Btw, the repo consistency tests will fail because the modular doesn't import EmuTextConfig. I found that the modular imports all the things specified in module-file import section + all things in old-model-file import section. But Emu3TextConfig is in neither of them, so prob we need also to check imports between one model files. I'll think how to fix that

[ArthurZucker]

You can tag @Cyrilvallez !

[Cyrilvallez]

Thanks a lot, great work! With the new modular version #34487, I think we can still improve a bit! Should be merged very soon, but this is already very nice imo if you don't want to wait 🤗

[Cyrilvallez]

I think this should be solved by #34487!

[ArthurZucker]

Waiting for the updates regarding @Cyrilvallez 's PR, will review again once updated

[qubvel]

heh, is something wrong with code owners?

[stevhliu]

Thanks :)

[zucchini-nlp]

Yeah, seems like it automatically tags all code owners depending on files touched/created...

@ArthurZucker, would be nice to not tag that many people at once

[ArthurZucker]

Nice thanks for iterating! My only comment is that I have not personnaly looked enough at the MIMI or the VQVAE from Chameleon you would know better, but the more standard the better!
A few nits but good to go IMO.

[ArthurZucker]

nice to have some expected outputs!

[zucchini-nlp]

Let's merge 🚀