Add Janus model
Author: yaswanth19
Pull request: yes
Labels: name: New model

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #35928 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@zucchini-nlp I’ve started working on Janus and would love to get some guidance. Right now, I’ve just created a skeleton and implemented the ImageProcessing class.

My first major hurdle is the CONFIG. The Janus config on the Hub is quite composite and non-standard. Standard values like `hidden_size`, `num_attention_heads`, etc., seem to be hardcoded in their implementation.

From a testing perspective, how should I approach writing the config class? Loading this config directly using AutoConfig.from_pretrained() doesn’t work.

Shall I write an ad hoc script to convert this config into a standard Hugging Face config (similar to convert_weights_to_hf.py, but for config)?

Config: https://huggingface.co/deepseek-ai/Janus-Pro-1B/blob/main/config.json

@yaswanth19 super cool to see a draft PR!

Yeah, that reminds me of Molmo is modeled, also with hardcoded values for configuration. I am not sure how you usually approach testing. If you will test by matching with the actual weights, then yes, converting a config will be helpful

As a first step, I'd suggest to make a working model code and then convert weights/config. The vision backbone should be very similar to existing CLIP models, and for VQ part feel free to look at Emu3 model [here](https://github.com/huggingface/transformers/tree/main/src/transformers/models/emu3). When the model is converted, we can try to match logits and see in which modules the logits start to diverge. LMK if you are stuck in any place and need help :)

Btw, I would be very interested to see if Janus can handle interleaved generation of image + text in one go 👀  If that's possible, would be super super nice

Hi @yaswanth19, this is really nice. I am also really interested in the model, do you think I can collaborate with you on this one? I can help with the implementation.

@geetu040 Thanks for your interest and for offering to collaborate! This is quite an ambitious PR for myself, and I’d like to take on the challenge of tackling it myself. That said, I’ll definitely reach out if I get stuck or unable to continue working on it :)

> @geetu040 Thanks for your interest and for offering to collaborate! This is quite an ambitious PR for myself, and I’d like to take on the challenge of tackling it myself. That said, I’ll definitely reach out if I get stuck or unable to continue working on it :)

Sure, understood, I wish you a very Good Luck.

Hello @yaswanth19. I understand that you want to do it by yourself. Yet, I'd like to insist on collaboration a bit further. I've spent a bit of time this week on reproducing their virtual environment and running their generative and understanding modes of Janus for 1B and 7B models on rented GPUs, as well as reading the original Janus paper, the SigLip paper (i.e. the "understanding" image encoder) and the LlamaGen paper (from which they take the generative image tokenizer). I was somewhat busy in this specific week, so did not have a lot of time for this. Still, I was about to start reading their code before trying to collaborate on integrating. As I've spent some time on this already, are you OK with trying to work together?

@hsilva664 Cool! since you have already made some progress IG we are on the same page. Let's start with implementing the Modular version. Janus has three components SigLip, VQ model, Llama - If I not wrong SigLip and Llama are staright forward with so I can pick those two and will create the structure and processor w.r.t. Now the middle part which is a VQ model ; you can start working on that. The reason I am insisting this way is that VQ model is a bit ambigous as I am not much familiar with that (Best case scenario we can inherit all the components for this model from existing ones :) ). Hence trying to code up the familiar part first and then iterate to match logits. WDYT?

That works. Also, on the Janus paper (the original one, not Janus pro), they state that they use DeepSeek-LLM (1.3B), not Llama, I think on Janus Pro they keep it that way, but change it to 7B. It is the VQ paper that uses Llama + the autoregressive decoder.

@zucchini-nlp I've added a preliminary version of the VQ part to [a clone](https://github.com/hsilva664/transformers/tree/yaswanth19-add-janus-model) of @yaswanth19 branch and am trying to set up a simple test code that will be used for checking weights/activations consistency with the original Janus model in the future (still have to import their weight values). From what I can gather, it seems that in HF the common practice is to have a JINJA template to convert between a list of dicts such as:

```
conversation = [
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": "What’s shown in this image?"},
            ],
    },
    {
        "role": "assistant",
        "content": [{"type": "text", "text": "This image shows a red stop sign."},]
    },
    {

        "role": "user",
        "content": [
            {"type": "text", "text": "Describe the image in more details."},
        ],
    },
]
```

To the processed string that will then be tokenized. This template seems to usually be stored on the hub, in JSON files like [this one](https://huggingface.co/llava-hf/llava-1.5-7b-hf/blob/main/chat_template.json). On their codebase, they have this conversion done by a Python module instead (see [here](https://github.com/deepseek-ai/Janus/blob/main/janus/utils/conversation.py), they use the template called "deepseek"). Just to confirm, should I change that format to a JINJA template, or would it be something "normal" to keep it like what they are doing (they call a function manually to convert the list inside the processor)? Was something similar to this done somewhere, so I can base myself on it (it is OK if not)?

@hsilva664 We have to create the jinja template. I have created one just for model testing purposes and it needs some corrections.

```
{
    "chat_template": "{% for message in messages %}{{ message['role'].capitalize() + ': '}}{% for content in message['content'] %}{% if content['type'] == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}{% if not loop.last %}\n\n{% endif %}{% endfor %}"
  }
```

My testing Conversation = [{"role":"User",
  "content":[{'type':"text","text":"<image_placeholder>\nConvert the formula into latex code.\n"}]},
  {"role": "Assistant", "content": " "},
]  

First of all this template doesn't account for new line chars before and after applying to conv. In the sense- to match the input_ids with original implementation we have to hackily use this way `system_prompt + '\n\n' + apply_chat_template(conv)+'\n'`. 
  

About the random initialization of weights (e.g. `EmuVQVAE.__init_weights__` or `Emu3PreTrainedModel.__init_weights__` in [this file](https://github.com/huggingface/transformers/blob/main/src/transformers/models/emu3/modular_emu3.py)). Are these meant to be used when the user wants to train from scratch (otherwise these would just be overwritten by the pretrained weights and serve no purpose...)? If so, unless I missed something, the Janus implementation does nothing like that and just leaves the default torch initialization. Should I do the same?

In the original implementation for Image generation they are generating a token and then applying `Classifier Free Guidance` after which processing it through some layers. But for text generation task it's using the HF implementation of `generate` function. AFAIK it won't be possible to have a single `generate` function for both these functionalities (Image+Text -> Text and Text->Image). 

Also I don't think it can perform Interleaved generation as mentioned by one of the authors [[LINK](https://huggingface.co/papers/2410.13848#67128c580f6cb5b1308e1528)]. For image generation they are explicitly adding <begin_of_image> tag to let the model know to generate image related logits. So I not sure how to add this image generation logic (In generate and prompt creation) without hampering the text generation functionality. One hack I can think of; is to use a flag where user explicitly mentions to generate image and we can process accordingly.

@zucchini-nlp Would like your thoughts on this issue and attached the image generation inference file below.

https://github.com/deepseek-ai/Janus/blob/1daa72fa409002d40931bd7b36a9280362469ead/generation_inference.py#L94

@hsilva664 

> Are these meant to be used when the user wants to train from scratch (otherwise these would just be overwritten by the pretrained weights and serve no purpose...)?

Yep, only for training from scratch and for some internal tests when a dummy model is init. If Janus uses a different initialization, you can add it in `_init_weights` 


@yaswanth19 really interesting point, didn't know Janus would add more layers after CFG 🙃 We have CFG in transformers, so it won't be a big problem. For the rest of layers we can override `generate()` and smth like below, WDYT?


```python
def generate(self, input_ids, **kwargs):
    # if generating text, just call super() 
    if kwargs.get("generation_mode", None) == "text":
        return super().generate(*input_ids, *kwargs)
    
    # Otherwise do custom logic, and rely only on greedy/sampling methods
    # Beam search, etc not supported

    # So some preprocessing, same as in `generation/utils.py`. Feel free to look at how MusicGen does it
    # Would have been nice to have self.prepare_for_generate(), planned and was low priority
    # I think since we have Janus planned, we can raise the priority a bit
    generation_config = ....

    for i in range(generation_config.max_new_tokens):
        next_token_outputs = super()._sample(input_ids, max_new_tokens=1 **kwargs)
        input_ids = self.layers_after_CFG(next_token_outputs)
    return next_token_outputs
```


cc @gante we're having first image generation model that doesn't work with `generate()` as is. I believe refactoring generation to tiny functions will help. We need to allow users to apply any pre/post-processing before and after the next token is sampled. So after separating out prefill and decoding, we can add an umbrella `self.get_next_token` maybe, though it will be hard to make it work with beam search and assisted decoding 😢 

@zucchini-nlp WDYT on image quality. Just wanted to confirm as you were a bit hesitant for ANOLE Model. These images are generated using original implementation itself (in worst case we can match it). The upper row images are generated using `Janus-Pro 7B` whereas the bottom row using `1B` variant.

**Prompt 1:** Generate an image of a snowman.
**Prompt 2:** A close-up high-contrast photo of Sydney Opera House sitting next to Eiffel tower, under a blue night sky of roiling energy, exploding yellow stars, and radiating swirls of blue.
 
<img width="762" alt="image" src="https://github.com/user-attachments/assets/dd2beb9a-a2c4-43ad-a39b-e64a7a2f75fc" />


@yaswanth19 yeah, looks good to me. The problem with Anole was that it would generate black screen or just some abstraction, instead of the image you asked. So as long as the requested image content in there, I'm oke :)

Btw, how long does it take to generate one image and how many tokens is that?

@zucchini-nlp I was comparing the generated images to SOTA diffusion models and was afraid the quality was not good 😅
> Btw, how long does it take to generate one image and how many tokens is that?

Each image is divided into 576 tokens and 1B variant took ~2 to 4 min in full precision locally whereas 7B took ~10 to 15 min

Hehe, nah, I don't expected autoregressive image generation to be better yet. Thanks, ~500 tokens isn't too bad. We can try to see how to get some speed ups later, 10-15 min is a lot (unless it was on CPU)

@zucchini-nlp One minor hiccup is how to get the `generation_mode` in modelling file. We have to use that argument in processor as depending on text/image generation we will format the prompt and use it in modelling to generate accordingly. Letting the User pass the same argument in `processor` and `generate` is not ideal. I can think of user passing that arg in Processor and we will be setting it  in Janusconfig. Then read that argument from Janus config. WDYT of the flow?

PS: For image generation we will add <begin_of_Image> tag at end letting model know to generate image tokens AFAIK.

@yaswanth19 I don't think it's a good idea to magically infer which generation mode is used when calling `generate()`, nor changing the model config file with processor. If the only difference in image-gen mode if the BOI token, then let's just delegate it to generation time by `ForceBOSToken` logits processor. Otherwise, imo it's oke to ask users to explicitly set generation mode and processing mode

Hello @zucchini-nlp . Me and @yaswanth19  just want to confirm something about the chat template. Each of us started doing it a different way and now we want to be all on the same page to adapt the code accordingly. Both of them work and it is just a matter of preference. The way I was doing it was the following: the template is passed in the following format (I'll already put interleaved images, as this is the most general case)

```
messages = [{
    "role": "user",
    "content": [
        {"type": "text", "text": "What’s the difference between "},
        {"type": "image"},
        {"type": "text", "text": " and "},
        {"type": "image"},
         {"type": "text", "text": "?"},
    ],
},
{
    "role": "assitant",
    "content": [
        {"type": "text", "text": "XXXXXX"}
    ],
},
...
]
```

Which the user would then process as follows:

```
text = processor.apply_chat_template(messages, add_generation_prompt=False)
```

Which--following the format in the Janus codebase--becomes:

```

: What’s the difference between 
<image_placeholder>
and 
<image_placeholder>


: XXXXXX<｜end▁of▁sentence｜>...
```

Then, text is to be passed to the tokenizer, then model... I did it like this mostly because I was reading Llava code and thought it was similar. @yaswanth19  had a chat template expecting messages as follows:

```
messages = [  
{  
"role": "User",  
"content": "What’s the difference between\n<image_placeholder>\nand\n<image_placeholder>\n?",  
},  
{"role": "Assistant", "content": "XXXXXX"},
...
]
```

What is the way we should stick to?

@zucchini-nlp  A few questions for some more clarity as we are almost ready and only a quater away(except tests 😅 ) .

1. Is Janus model and JanusForCasualLM class required, or can we only support `JanusForConditionalGeneration`, similar to PaliGemma?
2. Can you take a quick look at the `generate` function and help us with approaching/writing generation test cases?
* We want to run all test cases for text generation but only a subset for image generation.
   - How should we pass the `generation_mode` kwargs (fails validation and not required for other model to add it as default/base arg) and `generation_config` in tests?  
   - IMO we might have to overwrite most of image generation tests because we simply are not returning the next token but applying some pre and post processing.
* What all subset of tests are required for image generation?
3. What different test suites do we need to write? I can think of a test suite for the following:
   - Vision-to-Text
   - VQ Model
   - JanusModel if the class is created.
   - Text-to-Vision
   - Do we also need to support training for the VQ Module? Since in Anole, we are freezing the VQ module.  
4. Conceptually, my understanding is that for image generation, we can't use multiple prompts because we're using Classifier-Free Guidance (CFG).
   - Technically, we can, but it would degrade image quality as inputs would be diverse and images would deviate too much from promtps.  
   - So to generate multiple images where/how should we pass `num_images_to_generate` similar to num_return_sequences or let the user batch the same prompt to the required number?  


Hey @yaswanth19 @hsilva664 ! Glad to hear the PR is getting close to be finished :)

For the chat template, the first version with `{"type": "image"}` values is better suited and it will also support vectorized chat template outputs (see https://huggingface.co/docs/transformers/main/en/chat_template_multimodal). 

For the rest of questions:
1. We don't need a CausalLM class but a base model class would be nice and help us with the current vllm integration I am working on. The base model should be same as ConditionalGeneration but without an lm head, thus return hidden states 
2. Will leave comments below 👇🏻  For tests we will need to add basic `GenerationTesterMixin` for cases of generating text from images+text, all tests need to be ran and skipped if not supported. Since the default generation mode is `text`, no special actions needed. We have a plan to add model's default `GenerationConfig` class soon, which will help if we decide to add more test cases

For image generation we write our own tests, since it will be impossible to use the existing Mixin. Just a few test cases to make sure that the model runs, return correct output shape, type, can return dict etc will be enough. The generation mode is passed as part of `inputs_dict` or `generation_config`


3. Yes, the tests you listed will be needed. Not all of them are generative tests, for ex VQModel test will simply run all basic model tests and might need to skip a bunch. As an inspiration you can take a look at Anole VQ tests (https://github.com/zucchini-nlp/transformers/blob/db7a16d22bc93afc0af03c283c10faedfa4e7664/tests/models/anole/test_modeling_anole.py#L227-L241)

No we don't need to support VQ training. IMO that will add more complexity. We can consider adding a new model for VQ with its own modeling, if we see more models like Janus. For now let's leave as is

4. The best way is to allow users to batch inputs, same way as they are users to do with text. It is not clear to me, why CFG will be bad in batched generation? Btw, we have two CFG processor which you prob already noticed. IMO the one we have now in Janus is a bit weird to use, so what about if we use UnbatchedCFG, WDYT?



@yaswanth19 I'm not sure I follow 4 either. In the example Janus script from Deepseek codebase they take 1 prompt and then generate a lot of images for the same prompt (I think the default was 16). This agrees with tools such as LeonardoAI and Adobe AI, where you generate a lot of images for the same prompt and have the user choose one they think is best. This is just for obtaining diverse images and I don't think they inherently interfere with each other. That may be where the confusion came from, because, for a similar experience, we would need two nested loops: one iterating over prompts and another over images for the same prompt. As you know, CFG just means that we generate a condition-free logit and interpolate it with the conditional one before sampling, so I don't see how this would be a cause for an inability to use multiple prompts at the same time...

Got it w.r.t CFG. Misunderstanding from my side.

Btw, if you need several image per prompt, we can just set `num_return_sequences>1` in generation config, and expand inputs accordingly. Much easier than a double nested loop :)

@zucchini-nlp I'd like to report strange behaviour that I found when running some of the tests. In particular, for composite models, there seems to be inconsistent behaviour when `attn_implementation` is not passed explicitly on model creation. I have confirmed that this behaviour also happens on Emu3 and I believe it may happen in other models too. The problem goes away if `attn_implementation` is passed explicitly. I will describe it below.

The test creates a model using a simple `JanusModel(config)` call, where `config` is of type `JanusConfig` and contains `JanusTextConfig`, `JanusVisionConfig` and `JanusVQVAEConfig` as attributes. When running on a debugger, we can inspect that---for this same `JanusConfig`---the attribute `config._attn_implementation` starts with value `eager` before the call. On this model-creation call, however, a call to `PreTrainedModel.__init__()` happens via `super`, which then calls `PreTrainedModel._autoset_attn_implementation(...)`. This latter function, overrides `eager` with `sdpa`, which I believe might have been intentional. The problem is that the inner configurations do not get updated, and we end up with `config.vision_config` and `config.text_config` still set to `eager`. These latter configurations are the ones that get read by the attention layers when selecting their classes, so we end up with `eager` models and `sdpa` base config. The previously-existing tests in Siglip and also in the parent class `ModelTesterMixin` expect the NN to have sdpa layers if nothing gets passed, so our tests break if left unattended (I investigated the reason Siglip does not break, but it is beyond the point here). A hacky solution I wrote for now is just to do:

```python
if hasattr(self, "config"):  
    if hasattr(self.config, "vision_config"):  
        self.config.vision_config._attn_implementation = self.config._attn_implementation  
    if hasattr(self.config, "language_config"):  
        self.config.language_config._attn_implementation = self.config._attn_implementation
```

In `JanusPreTrainedModel` (i.e. force the other configurations to track the main one by hand right after it gets changed), but, of course, this is not a robust solution. I suspect this initialization behaviour is unintended.

@hsilva664 when initializing LM or vision backbones, you should use `_from_config(text_config)` which will not overwrite the config params. I agree that the attn implementation is not perfect and refactoring it is already planned. But init from config is a general way for multimodal models and should work I think 

@zucchini-nlp I don't understand exactly. Most tests (even beyond our new model) already have a line `model = model_class(config)`. It is this line that ends up calling `PreTrainedModel._autoset_attn_implementation(...)` that overwrites the configuration attention implementation. You mean all of them should have written `from_config()` instead...? 

Alternatively, did you mean that, inside the main model initialization, we should have initialized vision and language with `from_config`? If so, that would not solve the problem, for it is these same configs that become inconsistent after the call to `PreTrainedModel._autoset_attn_implementation(...)` (this happens in the very first line of model init, when calling `super().__init__`).

Furthermore, the existing tests also expect the attention implementation to be `sdpa`, so I believe overwriting `eager` with `sdpa` is supposed to happen. The problem is that it is happening for the model config only, but it should happen for its inner vision and language configs too.

For a quick reproduction of what I am saying you can do the following:
- run the test `Emu3Vision2TextModelTest::test_inputs_embeds_matches_input_ids`.
- place a breakpoint right before `model = model_class(config)`.
- Verify that `config._attn_implementation` is equal to `eager` and that `config.text_config._attn_implementation` is also equal to `eager`
- Run the line `model = model_class(config)`
- Verify that `config._attn_implementation` now changes to `sdpa`, but `config.text_config._attn_implementation` remains equal to `eager`

In out Janus model, this behaviour leads to the attention layers being instantiated as `eager` (because they look at vision/text configs, not base ones), but it should be `sdpa`. It was failing the test `test_sdpa_can_dispatch_composite_models` (now it is not failing anymore because I added the hack from the previous message where I make the inner configs follow the parent one by hand). The only reason Emu3 does not fail the test `test_sdpa_can_dispatch_composite_models` is because `Emu3Vision2TextModelTest` is marked with `_is_composite`  equal to False, so it skips the test. The only reason `SiglipVisionModelTest` passes this same test is because `id(model_sdpa.config) == id(model_sdpa.vision_model.config)` is True (i.e. they are the same object in memory), but that assertion should not be true for every model I think.

I'm leaving the manual `_attn_implementation` assignment for now, but please tell me if I should change it.

@hsilva664 that is weird, unless the way we set attention was broken recently. We did some major changes, so could have happened. I can take a look at it later. In the meanwhile, we had some tests for attention in test_sdpa_can_dispatch_composite_models and I see it passed for LLaVa on `main`, so you can try to write Janus same was as LLaVa

@zucchini-nlp  Upon inspection of Llava code, I see that they used `AutoModel.from_config(...)` in the main model initialization to create the inner language/vision models, so "option 2" from my previous message. It seems that this function creates new replicas of the text/vision configs and leaves the provided ones as lingering unused configurations. For clarity, see below, a snapshot taken from inside `LlavaForConditionalGeneration.__init__()` run from the test you mentioned:

```bash
id(config) = {int} 132268999633440
id(self.config) = {int} 132268999633440 # same as config
id(config.vision_config) = {int} 132268967891632
id(self.vision_tower.config) = {int} 132268915936864 # different
id(config.text_config) = {int} 132269125497296
id(self.language_model.config) = {int} 132268846369776 # different
```

Then, the solution would be to simply use `AutoConfig` in Janus and ignore the lingering `config.vision_config` and `config.text_config`. I think this will work for our VQ model. However, unfortunately, `JanusVisionTransformer`, our vision model, inherits from `SiglipVisionTransformer`, which does not have `PreTrainedModel` in its type hyerarchy. This seems to be necessary for creating with `from_config`. 

Looking more closely, I see that they have `SiglipModel` doing the following:

```python
# First, initialize the text and vision models with proper attention implementation
text_model = SiglipTextModel._from_config(text_config)
vision_model = SiglipVisionModel._from_config(vision_config)

# Second, get the text and vision submodules (for backward compatibility)
self.text_model = text_model.text_model
self.vision_model = vision_model.vision_model
```

So I guess we should try do the same. Instead of having the vision_model created as `JanusVisionTransformer`, we should first create a `tmp_vision_model` of type  `JanusVisionModel` containing `JanusVisionTransformer` and then do `self.vision_model = tmp_vision_model.vision_model`. What do you think?

@hsilva664 yep, that is exactly what I meant in the first message. If that works for Janus, then looks good to me. Adding more ways of setting attention is not a good practice, so let's try and re-use what we have in other VLMs :)

@zucchini-nlp The PR is ready for initial review. The tests CIs are not running currently but locally all the normal tests and Integration tests are passing; so good to go 🚀 .

PS: The tests are now running and passed.

@zucchini-nlp @yaswanth19 I think I've figured out the issue with the multi-prompt image generation. To contextualize @yaswanth19 reported that when using the Janus image generation mode and multiple different batched prompts, the model generates noise instead of images. The problem seems to be the padding tokens that get added when sentences of different sizes are combined. When I try to use a batch with two prompts of similar length, it generates almost correctly (sometimes there are errors), but if I instead try a very short prompt and then a very long one, the short prompt that gets left-padded has an output image that is simply noise, whereas the long prompt one has correct outputs. It is not the fact that there are multiple images present that causes this errors, it is simply the padding. In fact, I was able to reproduce the error in single prompt cases by artifically prepending the sequence with left padding before feeding it to the model.

There are some things that were happening when a shorter prompt was left-padded:

1. positions that correspond to the padding id were not being masked-out via attention mask and were therefore interferfering with the computation of outputs from attention layers. I believe this is unintended, since, if not for the mere fact that sequences have different lengths, these padding tokens would not even exist in the sequence. The function `GenerationMixin._prepare_attention_mask_for_generation` has routines for dealing with this, so I incorporated them. This reduced the number of noisy images significantly, but did not eliminate the problem. To quantify the improvement, I was running some tests where I was generating 20 images (script run 10 times, `generation_config.num_return_sequences` of 2) using a single left-padded artificial prompt that looks exactly like the left-padded does in the multi-prompt case. Prior to this modification, all 20 would come out as noise. This number was reduced to 7.
2. the positional embeddings will vary depending on the position of the token on the image sequence and this will be affected by padding. I checked how the Llama model was dealing with this for text. Correct if I am wrong, but it seems this difference is ignored there. The same sentence, under deterministic algorithms, will end up with different representations and logits depending on the padding and length of its co-occurring prompts. For Llama, this did not seem to be an issue and text generation works anyway, but it seems that Janus model is very sensitive to this. In particular, after some trial-and-error, I've concluded that what the model is really sensitive to is the position of the `bos` token. If the `bos` token occurs at the beggining of the prompt, the generation happens correctly. This is the case even if there are multiple padding tokens after the `bos` one. Therefore, I have added a routine that swaps the position of the leftmost padding token and the `bos` so that `bos` is always in the beginning (this routine can be made much simpler than my current code if desired by simply filling bos with padding and then adding bos at position zero, but I believe the way I wrote is more reliable). This significantly boosted multi-prompt generation. The 20 images test from before now yields zero noisy images.

I tested the two aforementioned changes for multi-image prompt and multi per-prompt images  in both 1B and 7B models and it seems to be working correctly. Swapping tokens like this is likely to be considered hacky and non-compliant with how other models do things. Still, the model seems to work with it. Do you agree with those changes?

> positions that correspond to the padding id were not being masked-out via attention mask 

Right, that exactly what has to be done. As a rule of thumb, we recommend to call `super.prepare_inputs_for_generation`, which is turn build positions ids using attention mask. So the only thing a user has to do is provide the mask, which is output from processor. I don't think we need to overwrite or anyhow replicate `_prepare_attention_mask_for_generation` as long as processor returns the correct mask

> I have added a routine that swaps the position of the leftmost padding token and the bos so that bos is always in the beginning

This is very hacky and undesired behavior. Not sure why it would be sensitive to the position of `bos`, if we are giving the correct position ids, and correct RoPE is applied to each embedding token. Can you check that positions are correctky passed, i.e. bos get position id `0` and so on? 

@zucchini-nlp thanks for your hints. I've updated the code and I think the current version complies better with expectations.

Indeed, there was a lingering attention mask on that function that had been passed by the processor (considering the "regular" script from the docs) and I was not using. I've changed the functions to make use of it and also added the call to `super().prepare_inputs_for_generation` together with including this mask in the previous pipeline where tensors are expanded (e.g. in cases of multiple images per prompt and because of CFG).

Also, the position ids were not being passed. I've changed the code to use them. I've had to modify some parts of the code. For example, the original Janus codebase had a line swapping all but the first token with `padding` that was also in our `generate` function, but we do not always want this, since bos will not be at position 0 if left padding occurs. There was also an issue where "position_ids" cannot be passed to `super().prepare_inputs_for_generation` because the necessary per-iteration position increments do not happen if we pass it (we have to pass None and let it infer the current position).

There is only one thing which you may find strange when you read the code after we merge it with this PR. As @yaswanth19 explained, we are using StaticCache here for the key and value caches. This is because we already know the size the final cache will need to have from the get-go, since we know the number of image tokens to be generated. Therefore, we do not need the computational burden of dynamically allocating more positions as needed. When the function `super().prepare_inputs_for_generation` gets called and detects we are using StaticCache, it converts the attention mask to 4d . There is even a comment saying this is necessary for  performant compiled forward pass. However, after the forward pass happens and the model generates logits, we have to call the function `self._update_model_kwargs_for_generation`. The implementation of this function already assumes that masks are 2d, not 4d (look at the way they concatenate the new token and this becomes clear). It errors if we pass 4d masks. I assume this is because the common scenarios where this function is used have `DynamicCache` instead of static. To not go into more detail, this required me to do some juggling with the masks to get the right ones to the right function. I think this is minor, but just letting you know.

@zucchini-nlp Adressed the comments and the PR is ready for review. Each image is taking around 35 sec on colab and I couldn't see any significant time diff when using static cache + torch.compile. There was 1-3 sec reduction when using cache+compilation but in few test runs it was matching that of a run without compilation so on avg not much difference

One surprising thing for me was with Dynamic cache the time taken was minimum around ~30 sec but with static cache it never decreased to that extent. May be I am missing some connection because of which we are not getting performance improvement with static cache+ compile🤔 

PS: Was running Janus-pro 1B on colab for these numbers.

Small code snippet I used for cache+compile.

```
from transformers import StaticCache, DynamicCache
from transformers import enable_full_determinism

enable_full_determinism(0)

# prompt_cache = DynamicCache()
model.generation_config.num_return_sequences = 1

model.generation_config.cache_implementation = "static"
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

outputs = model.generate(**inputs,
                        #  past_key_values = prompt_cache,
                         generation_mode="image",
                         do_sample=True,)
```

Thanks a lot! Yeah, the position ids are usually inferred from attention mask and we don't re-use user provided positions. It is in the general TODO list, but we should be good for Janus. It doesn't use special position ids

For static attention mask, we make it 4d to accomodate compile and avoid graph breaks indeed. AFAIK the `update_kwargs` method still should receive a 2D mask though, while the model get prepared inputs (4D mask, inferred position etc)

I will review the code in more details next week

@zucchini-nlp Addressed the comments 🤗 

@zucchini-nlp One last thing. I added the check you requested in the generate_image test, where we should compare the generated image with one uploaded to the HUB. There seems to be a problem where the result is hardware-dependent. On my machine, the test works correctly, but if I change the GPU to a remote one and compare with a file generated with my local hardware, there are differences in the image. The differences are not discernible when looking at the image, but they cause the assert to fail. Also, putting a small threshold on the assert is not enough, the difference reaches 255 in some channels of some pixels. The test already uses `sample=False`, but I tried additionally using  `enable_full_determinism`,  the issue is not gone. See the attached images, they correspond to the following 4 cases:

- local machine
- remote machine
- local with full determinism
- remote with full determinism

You can verify that their content is different using `sha256sum`. Should I remove this part from the test?

![img1](https://github.com/user-attachments/assets/d09cab2d-4dbd-4f03-bb8e-a449ad5990f2)
![img2](https://github.com/user-attachments/assets/9c68c621-8048-486b-8140-48fd4869d493)
![img3](https://github.com/user-attachments/assets/6a690388-684b-44b0-b5f2-682e9d32a828)
![img4](https://github.com/user-attachments/assets/17b68311-1dc1-4eaa-b5c8-d76b58302d75)

@hsilva664 yes, that is an issue we face in most other tests as well. I see the difference can be up to 255, but what if we compare bare pixels before postprocessing?

Also, I will help you before merge to match pixels with out runners, so we can have the tensor values for our hardware. Not a big deal for now, but thanks for flagging

@ArthurZucker DeepSeek is ready for final review. Note it has image-generation logic quite different so we override it, and support only greedy/sample generation. We'll see later if more models with special generation models are added in the library

@zucchini-nlp Also, when we run the model there is a warning saying

> The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.`

I did add `model.tie_weights()` in the convert script before saving, but the warning persists. Since the model works anyway, I ignored the warning, but just letting you know.

@ArthurZucker A gentle reminder for review whenever you have the bandwidth :hugs: 

@ArthurZucker A gentle ping for review 🤗 

Super sorry got caught up on Llama4 sprint

@ArthurZucker Addressed the comments !! 🤗 

Only two things are remaining before merging.
* Transfer the checkpoints to deepseek org space.
* Correct  the text_generate_images slow test (make it less flaky as its too depended on GPU ARCH) -> Probably compare the generated tokens instead of comparing decoded pixel values in the testcase. WDYT? @zucchini-nlp 

@yaswanth19 great! We can merge after the checkpoints are moved under DeepSeek org then. Yeah, checking tokens sounds good to me :)

Since we are modifying weights prob we need a new model repo for hf-janus. Do you have authors' contacts or if not, I can ask internally to transfer the weights. Which hub repo they are currently in?

@zucchini-nlp This is the `Janus` space in `deepseek-ai` org -> https://huggingface.co/collections/deepseek-ai/janus-6711d145e2b73d369adfd3cc

IG we can transfer the weights internally with HF prefix so the original weights are compatible with original Janus API.
Anyway tagging the Janus authors for their permission here @charlesCXK  @hills-code 

@yaswanth19 Thanks for your great work! I really appreciate your contribution and I am totally ok for the permission.

> IG we can transfer the weights internally with HF prefix so the original weights are compatible with original Janus API.

Great, sounds good! Ping me when PR to merge :)

@zucchini-nlp Can you please transfer the checkpoints to the [deepseek-ai](https://huggingface.co/collections/deepseek-ai/janus-6711d145e2b73d369adfd3cc) org.

1.) 1B ckpt: [yaswanthgali/Janus-Pro-1B-HF](https://huggingface.co/yaswanthgali/Janus-Pro-1B-HF)  
2.) 7B ckpt: [hugosilva664/Janus-Pro-7B-HF](https://huggingface.co/hugosilva664/Janus-Pro-7B-HF)


Once done will change the model_id in PR and should be good to go :rocket:

I guess you waned to ping @hills-code for weights :)

MB :sweat_smile: , @hills-code Can you please transfer the checkpoints to `deepseek-ai` in HUB or could you tag someone who can help with that?

@yaswanth19 I have updated with @charlesCXK and shared our conversations. I may not have the permission as I have checked out from DeepSeek

@yaswanth19 
Thank you so much for your incredible contributions and the outstanding work you've done. I truly appreciate the effort put into this project.

I wanted to share a minor concern regarding the hosting location: Uploading this to the `deepseek-ai` organization might not be the most suitable approach at this stage. Given the community's heightened attention to official updates in this space, there's a possibility it could be misinterpreted as an official model release from `deepseek-ai`, which might inadvertently generate unintended discussions.

May I suggest keeping the current Hugging Face repository URL as the hosting solution? This approach would allow us to maintain clarity while still showcasing your valuable work effectively 🌹.

@charlesCXK I see, in that case we can host models is a new org like `DeepSeek-community`. We're also wrapping up shipping DeepSeek-VL2 soon, I believe it can also be hosted under the community org. 

@yaswanth19 feel free to create an org and port checkpoints there, with some readmes and demo inference scripts :)

@zucchini-nlp The PR is good for merge. The failing tests are unrelated to this PR 🚀 🤗 

Here is the deepseek-community org and both the 1B and 7B are uploaded there. 

LINK: https://huggingface.co/deepseek-community

Super super cool, thanks for adding an org! Let's resolve merge conflicts and I'll merge  

@zucchini-nlp Resolved the Merge Conflict and lets merge it before another one appears 😅 

Let's go 🚀 

Hi @yaswanth19 @zucchini-nlp @hsilva664 , thanks for your excellent work! While using Janus for text to image generation, I met the following issue:
"JanusForConditionalGeneration has no `_prepare_4d_causal_attention_mask_with_cache_position` method defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're writing code, see Llama for an example implementation. If you're a user, please report this issue on GitHub."

Any insights or guidance on this would be greatly appreciated. 

> Hi @yaswanth19 @zucchini-nlp @hsilva664 , thanks for your excellent work! While using Janus for text to image generation, I met the following issue: "JanusForConditionalGeneration has no `_prepare_4d_causal_attention_mask_with_cache_position` method defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're writing code, see Llama for an example implementation. If you're a user, please report this issue on GitHub."
> 
> Any insights or guidance on this would be greatly appreciated.

To provide more context, I used the exact text to image generation code example from https://huggingface.co/docs/transformers/main/en/model_doc/janus, and this issue occurred while calling ``model.generate()``.

> > Hi @yaswanth19 @zucchini-nlp @hsilva664 , thanks for your excellent work! While using Janus for text to image generation, I met the following issue: "JanusForConditionalGeneration has no `_prepare_4d_causal_attention_mask_with_cache_position` method defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're writing code, see Llama for an example implementation. If you're a user, please report this issue on GitHub."
> > Any insights or guidance on this would be greatly appreciated.
> 
> To provide more context, I used the exact text to image generation code example from https://huggingface.co/docs/transformers/main/en/model_doc/janus, and this issue occurred while calling `model.generate()`.

My transformers version is 4.52.4.

I think I know where the issue stems from. @yaswanth19 want to submit a PR or I can do it as well? We can just add `_prepare_4d_causal_attention_mask_with_cache_position` or another options is to apply the new attn mask API (as in llama) with `create_causal_mask`

Thanks @Dswn888 for reporting the issue. I don't have an idea rn on why this resurfaced but will look into it. Yep @zucchini-nlp , I raise a PR, with new masking API to be up to date :thinking: 

Thanks for your prompt reply @yaswanth19  @zucchini-nlp ! 
By the way, could you please explain the potential impact of this issue?

Hey @Dswn888 , I ran the code with specified transformers version but couldn't replicate the issue. Even on main this doesn't exist 😢 . Can you run `transformers-cli env` and share the complete output. Also LMK if this issue persists on main in your setup. I am running this [example](https://huggingface.co/docs/transformers/main/en/model_doc/janus#single-image-inference)

Hi @yaswanth19 @zucchini-nlp @geetu040 @hsilva664 @ArthurZucker , 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**?

> Hey @Dswn888 , I ran the code with specified transformers version but couldn't replicate the issue. Even on main this doesn't exist 😢 . Can you run `transformers-cli env` and share the complete output. Also LMK if this issue persists on main in your setup. I am running this [example](https://huggingface.co/docs/transformers/main/en/model_doc/janus#single-image-inference)

Hi! Appreciate you trying to replicate the issue.

Here are my environment details from `transformers-cli env`:

- `transformers` version: 4.52.4
- Platform: Linux-4.19.90-2107.6.0.0192.8.oe1.bclinux.x86_64-x86_64-with-glibc2.35
- Python version: 3.11.11
- Huggingface_hub version: 0.32.3
- Safetensors version: 0.5.3
- Accelerate version: 1.7.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.7.0+cu126 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A800 80GB PCIe

>I hope to confirm with the official whether the image generation part of Janus supports multimodal input?

Hey @Dswn888  ,Yep, on image generation it's a simple T2I and doesn't support multi-modal inputs or output. And regarding the issue - Are you still facing the issue :thinking: try later versions or from main as I am unable to build the specified env in my setup due to cuda issues :sweat_smile: 


> > I hope to confirm with the official whether the image generation part of Janus supports multimodal input?
> 
> Hey @Dswn888 ,Yep, on image generation it's a simple T2I and doesn't support multi-modal inputs or output. And regarding the issue - Are you still facing the issue 🤔 try later versions or from main as I am unable to build the specified env in my setup due to cuda issues 😅

Hey, thanks for the quick reply and confirming the T2I aspect!

I was wondering a couple of things building on that:

1.  Did the original DeepSeek official Janus model support multimodal inputs for image generation?
2.  If I were to adapt the code to handle multimodal input, would the current model weights be able to understand the image-based information, or would SFT be necessary for that?

Also, about the issue we discussed earlier, it's still popping up on my end. I'll definitely try pulling a later version or from `main` to see if that fixes it.

> Did the original DeepSeek official Janus model support multimodal inputs for image generation.

AFAIK, No if it was supported back then, I would have added that feature. So, ig it was not supported

>If I were to adapt the code to handle multimodal input, would the current model weights be able to understand 

I had thought about it and tried some thing similar to image editing where I pass the image and supporting text in the hope of getting a edited image but was getting OOM error coz the input would be dense (576 image tokens + n text tokens) and on top of that we have generate 576 tokens for the image. If you have enough ram for 1B variant then it's worth a try. 

To make what i said to work you might need to do some hacks:
* Process the text + image in the same manner as we process for text generation (I would prefer not to use chat template here)
* Next generate the embeddings for the same and If I recall correctly then you have to use `get_image_features` on your end to adjust the image embeddings. 
* Next you can pass the input embeds which now contains the text and image info to model (comment out the image embeddings adjustment part) to generate the image tokens. some hacks here and there and it should work.

> > Did the original DeepSeek official Janus model support multimodal inputs for image generation.
> 
> AFAIK, No if it was supported back then, I would have added that feature. So, ig it was not supported
> 
> > If I were to adapt the code to handle multimodal input, would the current model weights be able to understand
> 
> I had thought about it and tried some thing similar to image editing where I pass the image and supporting text in the hope of getting a edited image but was getting OOM error coz the input would be dense (576 image tokens + n text tokens) and on top of that we have generate 576 tokens for the image. If you have enough ram for 1B variant then it's worth a try.
> 
> To make what i said to work you might need to do some hacks:
> 
> * Process the text + image in the same manner as we process for text generation (I would prefer not to use chat template here)
> * Next generate the embeddings for the same and If I recall correctly then you have to use `get_image_features` on your end to adjust the image embeddings.
> * Next you can pass the input embeds which now contains the text and image info to model (comment out the image embeddings adjustment part) to generate the image tokens. some hacks here and there and it should work.

Thanks for the detailed explanation and the insights! I really appreciate you taking the time to explain the challenges with multimodal inputs and the suggested hacks.

I have two follow-up questions if you don't mind:

1. Regarding processing text and images, you mentioned preferring not to use a chat template. The current chat template looks like this:

> {
>     "role": "user",
>     "content": [
>             {"type": "text", "text": "What’s the difference between"},
>             {"type": "image", "url": image_urls[0]},
>             {"type": "text", "text": " and "},
>             {"type": "image", "url": image_urls[1]}
>     ]
> }

Could you elaborate a bit on **why you'd advise against using this type of chat template** in this specific context?

2. Would you be open to **sharing the modified code or any snippets from your earlier attempts at image editing** where you were trying to pass images and supporting text? It would be incredibly helpful for my understanding and further experimentation.

Thanks again for all your help!

> Could you elaborate a bit on why you'd advise against using this type of chat template in this specific context?

Sorry MB, we can utilize chat_template - I was of the opinion the we only utilize chat_template for text generation but we also utilize for image gen so no issues and infact better as model is trained that way.

Regarding the scripts I unfortunately don't have those code snippets. Just to reiterate conceptually what I was saying above is we can pass multi modal inputs for text generation so we want the processing of text gen pipeline (so to get multimodal inputs) and generate of image pipeline (so for image gen pipeline this is like a long text and hopefully it identifies the multimodal info)
what ever hacks I have suggested above were just to fit these two halves of pipeline.

> > Could you elaborate a bit on why you'd advise against using this type of chat template in this specific context?
> 
> Sorry MB, we can utilize chat_template - I was of the opinion the we only utilize chat_template for text generation but we also utilize for image gen so no issues and infact better as model is trained that way.
> 
> Regarding the scripts I unfortunately don't have those code snippets. Just to reiterate conceptually what I was saying above is we can pass multi modal inputs for text generation so we want the processing of text gen pipeline (so to get multimodal inputs) and generate of image pipeline (so for image gen pipeline this is like a long text and hopefully it identifies the multimodal info) what ever hacks I have suggested above were just to fit these two halves of pipeline.

Got it, thanks so much for the detailed explanations!