Emu3: add model
Author: zucchini-nlp
Pull request: yes
Labels: name: New model

name: Multimodal

# 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:
- [x] Batched generation
- [x] Upload chat template and change the image-placeholder token from `extra-0` to smth like `<image>`
- [x] Match the orig implementation on logit level
- [x] Tests, many more tests
- [x] 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

```python
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: 
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:

```python
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("
", return_tensors="pt")[0] # torch.tensor([processor.tokenizer.eos_token_id], device=model.device)
    pad_token_id = processor.tokenizer.encode("
", return_tensors="pt")[0] # torch.tensor([processor.tokenizer.pad_token_id], device=model.device)
    eol_token_id = processor.tokenizer.encode("
", return_tensors="pt")[0]
    eof_token_id = processor.tokenizer.encode("
", 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")

```


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33770). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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 

You can tag @Cyrilvallez !

Yes, lets' wait for the new modular and then come back to this model, it is not very urgent

should be ready for the final review. I rebased on `main` with the latest modular refactor that works better for Emu3

heh, is something wrong with code owners?

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

Let's merge 🚀