Repository: yaswanth19/transformers
Pull request #6: add-janus-model
Author: login: hsilva664
Created: 2025-03-07T23:17:18Z | Merged: 2025-03-08T17:13:48Z

# Nex fixes

- Adding JanusVisionModel as a proper wrapper
- Fixing multi-prompt image generation
- All tests have passed, but we still need to update the files in the HUB to reflect the new changes


Diff:
diff --git a/src/transformers/models/janus/configuration_janus.py b/src/transformers/models/janus/configuration_janus.py
index 90c5190c2825..433ce752340c 100644
--- a/src/transformers/models/janus/configuration_janus.py
+++ b/src/transformers/models/janus/configuration_janus.py
@@ -43,6 +43,8 @@ class JanusVisionConfig(PretrainedConfig):
     Args:
         hidden_size (`int`, *optional*, defaults to 1024):
             Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to `hidden_size * mlp_ratio`):
+            Dimensionality of the intermediate layer in the MLP block.
         num_hidden_layers (`int`, *optional*, defaults to 24):
             Number of hidden layers in the Transformer encoder.
         num_attention_heads (`int`, *optional*, defaults to 16):
diff --git a/src/transformers/models/janus/convert_janus_weights_to_hf.py b/src/transformers/models/janus/convert_janus_weights_to_hf.py
index 30df02de66ef..fd6b0d3db0ee 100644
--- a/src/transformers/models/janus/convert_janus_weights_to_hf.py
+++ b/src/transformers/models/janus/convert_janus_weights_to_hf.py
@@ -45,11 +45,11 @@
 # Mappings
 MAPPINGS = {
     # Vision model
-    r"(?<!gen_)vision_model.vision_tower.blocks": "model.vision_model.encoder.layers",
-    r"(?<!gen_)vision_model.vision_tower.pos_embed": "model.vision_model.embeddings.position_embeddings.weight",
-    r"(?<!gen_)vision_model.vision_tower.patch_embed.proj": "model.vision_model.embeddings.patch_embeddings.projection",
-    r"(?<!gen_)vision_model.vision_tower.norm": "model.vision_model.post_layernorm",
-    r"(?<!gen_)vision_model.vision_tower.attn_pool": "model.vision_model.head",
+    r"(?<!gen_)vision_model.vision_tower.blocks": "model.vision_model.vision_model.encoder.layers",
+    r"(?<!gen_)vision_model.vision_tower.pos_embed": "model.vision_model.vision_model.embeddings.position_embeddings.weight",
+    r"(?<!gen_)vision_model.vision_tower.patch_embed.proj": "model.vision_model.vision_model.embeddings.patch_embeddings.projection",
+    r"(?<!gen_)vision_model.vision_tower.norm": "model.vision_model.vision_model.post_layernorm",
+    r"(?<!gen_)vision_model.vision_tower.attn_pool": "model.vision_model.vision_model.head",
     r"(?P<pre>\b(vision_model|model\.vision_model)\b.*\.)proj(?=\.|\s|$)": r"\g<pre>projection_layer",
     r"(?P<pre>\b(vision_model|model\.vision_model)\b.*\.)norm(?=\.|\s|$)": r"\g<pre>layer_norm",
     r"(?P<pre>\b(vision_model|model\.vision_model)\b.*\.)norm1(?=\.|\s|$)": r"\g<pre>layer_norm1",
@@ -157,9 +157,10 @@ def convert_state_dict_to_hf(state_dict):
                 converted_state_dict[new_key] = state_dict[old_key]
 
     # Embeddings will not have initial dimension
-    converted_state_dict["model.vision_model.embeddings.position_embeddings.weight"] = converted_state_dict[
-        "model.vision_model.embeddings.position_embeddings.weight"
-    ].squeeze(0)
+    converted_state_dict["model.vision_model.vision_model.embeddings.position_embeddings.weight"] = (
+        converted_state_dict[
+            "model.vision_model.vision_model.embeddings.position_embeddings.weight"
+        ].squeeze(0))
 
     return converted_state_dict
 
@@ -404,7 +405,8 @@ def convert_model(
     model.generation_config.guidance_scale = 5
     model.generation_config.pad_token_id = tokenizer.vocab.get(
         "<\uff5c\u2581pad\u2581\uff5c>"
-    )  # If fixed for all janus variants then hardcode it.
+    )
+    model.generation_config.generation_kwargs["boi_token_id"] = tokenizer.vocab.get("<begin_of_image>")
 
     # Load and convert state dict
     print("Loading state dict...")
diff --git a/src/transformers/models/janus/modeling_janus.py b/src/transformers/models/janus/modeling_janus.py
index a44915ef918e..35d386495cc6 100644
--- a/src/transformers/models/janus/modeling_janus.py
+++ b/src/transformers/models/janus/modeling_janus.py
@@ -1346,8 +1346,7 @@ def __init__(self, config: JanusConfig):
         super().__init__(config)
         self.config = config
         # This is necessary for backward compatibility, see SiglipModel initialization
-        tmp_vision_model = JanusVisionModel._from_config(config.vision_config)
-        self.vision_model = tmp_vision_model.vision_model
+        self.vision_model = JanusVisionModel._from_config(config.vision_config)
         self.aligner = JanusVisionAlignerMLP(self.vision_model.config)
 
         self.vqmodel = JanusVQVAE._from_config(config.vq_config)
@@ -1375,6 +1374,9 @@ def get_image_embeddings(self, pixel_values):
         image_embeds = self.aligner(image_embeds.last_hidden_state)
         return image_embeds
 
+    def _prepare_4d_causal_attention_mask_with_cache_position(self, *args, **kwargs):
+        return self.language_model._prepare_4d_causal_attention_mask_with_cache_position(*args, **kwargs)
+
     @add_start_docstrings_to_model_forward(JANUS_INPUTS_DOCSTRING)
     def forward(
         self,
@@ -1682,6 +1684,7 @@ def generate(
         # 6. Expand inputs for multiple image generations per prompt.
         input_ids, model_kwargs = self._expand_inputs_for_generation(
             input_ids=input_ids,
+            attention_mask=attention_mask,
             expand_size=generation_config.num_return_sequences,
             **model_kwargs,
         )
@@ -1691,9 +1694,17 @@ def generate(
 
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
-        # Should we double nly when guidance scale is not None.
+        # Should we double only when guidance scale is not None.
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
-        input_tokens[batch_size:, 1:-1] = generation_config.pad_token_id  # Set Unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None)
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(
+            (input_tokens[batch_size:, :] != generation_config.bos_token_id)
+            & (input_tokens[batch_size:, :] != generation_config.generation_kwargs["boi_token_id"]),
+            generation_config.pad_token_id,
+        )
+
         inputs_embeds = self.get_input_embeddings()(input_tokens)
 
         model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
@@ -1702,7 +1713,7 @@ def generate(
             # Prepare cache if not provided
             model_kwargs["past_key_values"] = self._get_cache(
                 cache_implementation=generation_config.cache_implementation or "static",
-                # batch_size should account for both conditional/unconditional input; hence mulitplied by 2.
+                # batch_size should account for both conditional/unconditional input; hence multiplied by 2.
                 batch_size=batch_size * 2,
                 # we should have at least a cache len of seq_len + num_image_tokens.
                 max_cache_len=max(generation_config.max_length, num_image_tokens + seq_len),
@@ -1727,17 +1738,24 @@ def generate(
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
         for i in range(num_image_tokens):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
+            model_inputs = self.prepare_inputs_for_generation(
+                inputs_embeds=inputs_embeds, input_ids=input_tokens, attention_mask=attention_mask, **model_kwargs
+            )
+
+            # inputs_embeds.device can change on multi-gpu
+            model_inputs["attention_mask"] = model_inputs["attention_mask"].to(inputs_embeds.device)
+            model_inputs["cache_position"] = model_inputs["cache_position"].to(inputs_embeds.device)
+
             outputs = self.model.language_model(
-                inputs_embeds=inputs_embeds,
                 output_attentions=output_attentions,
                 output_hidden_states=output_hidden_states,
-                **model_kwargs,
+                **model_inputs,
             )
 
             # Update model_kwargs like cache_position for next generation.
+            model_kwargs["attention_mask"] = attention_mask  # needed for the following update
             model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+            attention_mask = model_kwargs.pop("attention_mask", None)  # to avoid future in-place modification
             hidden_state = outputs.last_hidden_state[:, -1, :].clone()
 
             # Generate scores using the generation head. (not using above defined lm head)
@@ -1759,6 +1777,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = next_token[:, None]
+
         if return_dict_in_generate:
             if output_scores:
                 scores += (scores,)
diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py
index 067571626db1..34c1802eab17 100644
--- a/src/transformers/models/janus/modular_janus.py
+++ b/src/transformers/models/janus/modular_janus.py
@@ -148,6 +148,7 @@ def _init_weights(self, module):
                 module.weight.data[module.padding_idx].zero_()
 
 
+
 @dataclass
 class JanusVQVAEOutput(ModelOutput):
     """
@@ -881,8 +882,7 @@ def __init__(self, config: JanusConfig):
         super().__init__(config)
         self.config = config
         # This is necessary for backward compatibility, see SiglipModel initialization
-        tmp_vision_model = JanusVisionModel._from_config(config.vision_config)
-        self.vision_model = tmp_vision_model.vision_model
+        self.vision_model = JanusVisionModel._from_config(config.vision_config)
         self.aligner = JanusVisionAlignerMLP(self.vision_model.config)
 
         self.vqmodel = JanusVQVAE._from_config(config.vq_config)
@@ -910,6 +910,10 @@ def get_image_embeddings(self, pixel_values):
         image_embeds = self.aligner(image_embeds.last_hidden_state)
         return image_embeds
 
+    def _prepare_4d_causal_attention_mask_with_cache_position(self, *args, **kwargs):
+        return self.language_model._prepare_4d_causal_attention_mask_with_cache_position(*args, **kwargs)
+
+
     @add_start_docstrings_to_model_forward(JANUS_INPUTS_DOCSTRING)
     def forward(
         self,
@@ -1217,6 +1221,7 @@ def generate(
         # 6. Expand inputs for multiple image generations per prompt.
         input_ids, model_kwargs = self._expand_inputs_for_generation(
             input_ids=input_ids,
+            attention_mask=attention_mask,
             expand_size=generation_config.num_return_sequences,
             **model_kwargs,
         )
@@ -1226,9 +1231,17 @@ def generate(
 
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
-        # Should we double nly when guidance scale is not None.
+        # Should we double only when guidance scale is not None.
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
-        input_tokens[batch_size:, 1:-1] = generation_config.pad_token_id  # Set Unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None)
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(
+            (input_tokens[batch_size:, :] != generation_config.bos_token_id) &
+            (input_tokens[batch_size:, :] != generation_config.generation_kwargs["boi_token_id"]),
+            generation_config.pad_token_id
+        )
+
         inputs_embeds = self.get_input_embeddings()(input_tokens)
 
         model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
@@ -1237,7 +1250,7 @@ def generate(
             # Prepare cache if not provided
             model_kwargs["past_key_values"] = self._get_cache(
                 cache_implementation=generation_config.cache_implementation or "static",
-                # batch_size should account for both conditional/unconditional input; hence mulitplied by 2.
+                # batch_size should account for both conditional/unconditional input; hence multiplied by 2.
                 batch_size=batch_size * 2,
                 # we should have at least a cache len of seq_len + num_image_tokens.
                 max_cache_len=max(generation_config.max_length, num_image_tokens + seq_len),
@@ -1261,18 +1274,29 @@ def generate(
         decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
+
         for i in range(num_image_tokens):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
-            outputs = self.model.language_model(
+            model_inputs = self.prepare_inputs_for_generation(
                 inputs_embeds=inputs_embeds,
+                input_ids=input_tokens,
+                attention_mask=attention_mask,
+                **model_kwargs
+            )
+
+            # inputs_embeds.device can change on multi-gpu
+            model_inputs["attention_mask"] = model_inputs["attention_mask"].to(inputs_embeds.device)
+            model_inputs["cache_position"] = model_inputs["cache_position"].to(inputs_embeds.device)
+
+            outputs = self.model.language_model(
                 output_attentions=output_attentions,
                 output_hidden_states=output_hidden_states,
-                **model_kwargs,
+                **model_inputs,
             )
 
             # Update model_kwargs like cache_position for next generation.
+            model_kwargs["attention_mask"] = attention_mask  # needed for the following update
             model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+            attention_mask = model_kwargs.pop("attention_mask", None) # to avoid future in-place modification
             hidden_state = outputs.last_hidden_state[:, -1, :].clone()
 
             # Generate scores using the generation head. (not using above defined lm head)
@@ -1294,6 +1318,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = next_token[:, None]
+
         if return_dict_in_generate:
             if output_scores:
                 scores += (scores,)
diff --git a/tests/models/janus/test_modeling_janus.py b/tests/models/janus/test_modeling_janus.py
index f804cd3f1ad9..f580171a7d43 100644
--- a/tests/models/janus/test_modeling_janus.py
+++ b/tests/models/janus/test_modeling_janus.py
@@ -17,7 +17,7 @@
 import tempfile
 import unittest
 from functools import reduce
-
+import re
 import numpy as np
 import requests
 
@@ -279,17 +279,13 @@ def test_sdpa_can_dispatch_composite_models(self):
 
             for name, submodule in model_eager.named_modules():
                 class_name = submodule.__class__.__name__
-                if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name:
-                    raise ValueError("The eager model should not have SDPA attention layers")
+                if any(re.finditer(r"Attention(?!Pool)", class_name)):
+                    self.assertTrue(submodule.config._attn_implementation == "eager")
 
-            has_sdpa = False
             for name, submodule in model_sdpa.named_modules():
                 class_name = submodule.__class__.__name__
-                if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name:
-                    has_sdpa = True
-                    break
-            if not has_sdpa and model_sdpa.config.model_type != "falcon":
-                raise ValueError("The SDPA model should have SDPA attention layers")
+                if any(re.finditer(r"Attention(?!Pool)", class_name)):
+                    self.assertTrue(submodule.config._attn_implementation == "sdpa")
 
     def check_training_gradient_checkpointing(self, gradient_checkpointing_kwargs=None):
         if not self.model_tester.is_training:
@@ -301,7 +297,7 @@ def check_training_gradient_checkpointing(self, gradient_checkpointing_kwargs=No
         that are matched with images on Siglip, as the forward pass of our models---which is what is used to check
         gradients in this test---does not include the forward pass on this layer.
         """
-        skip_patterns = ["vision_model.head", "vqmodel", "gen_embed", "gen_aligner", "gen_head"]
+        skip_patterns = ["vision_model.vision_model.head", "vqmodel", "gen_embed", "gen_aligner", "gen_head"]
 
         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__):


Reviews:
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T05:59:03Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:18:04Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:20:01Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T10:26:36Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:35:37Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:41:21Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:45:09Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:47:23Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:50:18Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T10:51:39Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T10:54:32Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T10:54:35Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:56:27Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T10:57:56Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T10:59:21Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T11:01:18Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T11:06:10Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T11:10:18Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T11:19:11Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T11:25:34Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T11:26:27Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T13:19:27Z
user: login: hsilva664
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-08T13:25:30Z
user: login: yaswanth19
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-08T13:39:31Z

Review comments:
diff_hunk: @@ -1119,8 +1122,7 @@ def __init__(self, config: JanusConfig):
         super().__init__(config)
         self.config = config
         # This is necessary for backward compatibility, see SiglipModel initialization
-        tmp_vision_model = JanusVisionModel._from_config(config.vision_config)
-        self.vision_model = tmp_vision_model.vision_model
+        self.vision_model = JanusVisionModel._from_config(config.vision_config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: If we are creating the vision model then is there a need for the overwriting sdpa tests ig. IMO we should be able to do the below code change and this should fix set the attn_implementation.

`vision_model = JanusVisionModel._from_config(config)` 
created_at: 2025-03-08T02:48:14Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: This one feels odd when considering repo consistency, lets instantiate the the vision model and since it's a wrapper around vision transformers so no harm in logits and keep this code design intact. Also vision transformers is not common vision backbone name we see, we generally keep it as vision_tower, vision_model etc.
```suggestion
        self.vision_model = JanusVisionModel(config)
```
created_at: 2025-03-08T02:51:45Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
         self.post_init()
 
+    def forward(self, *args, **kwargs):
+        return self.vision_transformer(*args, **kwargs)
+
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: We can't use this as its too messy (don't know which argument is getting passed in code flow). Pass the standard args `pixel_values,return_dict,output_attention etc.. and remaining as kwargs`. Infact this seems to be a standard vsiion model module, so check can we use modular to import this module from another model
created_at: 2025-03-08T03:34:48Z
author_association: OWNER
diff_hunk: @@ -1390,7 +1398,7 @@ def prepare_inputs_for_generation(
                 batch_size, sequence_length = input_ids.shape
                 device = input_ids.device
 
-            attention_mask = self.model.language_model._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: Do we need this change? With this change we are adding an extra function.
created_at: 2025-03-08T03:38:21Z
author_association: OWNER
diff_hunk: @@ -1544,18 +1557,29 @@ def generate(
         decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
+
         for i in tqdm(range(num_image_tokens)):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
-            outputs = self.model.language_model(
+            model_inputs = super().prepare_inputs_for_generation(
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: This is not the correct implementation, we already have a prepare_input_for_generation function defined above and there we have to do this super call. I have made the changes and will push it in sometime
created_at: 2025-03-08T03:44:23Z
author_association: OWNER
diff_hunk: @@ -1510,8 +1519,12 @@ def generate(
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None) 
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(input_tokens[batch_size:, :] != generation_config.bos_token_id,
+                                                  generation_config.pad_token_id)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: Need some more clarification - I think in the original implementation they are masking every token except first and last which would be `beign_of_sentence_token` and `begin_of_image_token`(bos,boi). and here we are doing tokens !=bos which would also mask boi token
created_at: 2025-03-08T03:54:55Z
author_association: OWNER
diff_hunk: @@ -1572,6 +1596,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = torch.cat([input_tokens, next_token[:, None]], dim=-1)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: Can you clarify why this is required, AFAIK input embeds are used for first generation step and later on the next generated token is used along with cache. so we process next token and for previous context use the cache but here we are concating the input_tokens with next token, IMO we are passing the entire context everytime.
created_at: 2025-03-08T05:58:46Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I changed the name because otherwise the name of the variable would be `model.vision_model.vision_model` in our `state_dict`, which I think is not desired. Siglip does it like this, but, again, siglip uses that `tmp` construct that we've agreed to change (i.e., create a `VisionModel` and then passes its inner `vision_model` attribute and the parent gets deleted). What name do you want then?
created_at: 2025-03-08T10:18:04Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
         self.post_init()
 
+    def forward(self, *args, **kwargs):
+        return self.vision_transformer(*args, **kwargs)
+
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: Will check once we decide on the name of the argument. Will either add the explicit arguments like you are suggesting or leave these lines out depending on your decision.
created_at: 2025-03-08T10:20:00Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: Ohh yeah then we can keep `model = JanusVisionModel()`. Then we will have model.vision_model.model.vision_transformer. Please refer to `siglip2` codebase. The problem with siglip is they have used tmp for BC of previous checkpoints as it was integrated a few years ago. 
created_at: 2025-03-08T10:26:36Z
author_association: OWNER
diff_hunk: @@ -1390,7 +1398,7 @@ def prepare_inputs_for_generation(
                 batch_size, sequence_length = input_ids.shape
                 device = input_ids.device
 
-            attention_mask = self.model.language_model._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: After the other changes, `generate` was trying to call (under the hood) `_prepare_4d_causal_attention_mask_with_cache_position` right from `JanusForConditionalGeneration` in one of the functions below and throwing a warning because it did not have it.
created_at: 2025-03-08T10:35:36Z
author_association: NONE
diff_hunk: @@ -1119,8 +1122,7 @@ def __init__(self, config: JanusConfig):
         super().__init__(config)
         self.config = config
         # This is necessary for backward compatibility, see SiglipModel initialization
-        tmp_vision_model = JanusVisionModel._from_config(config.vision_config)
-        self.vision_model = tmp_vision_model.vision_model
+        self.vision_model = JanusVisionModel._from_config(config.vision_config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: `sdpa` parent tests---which you are suggesting reverting to---try to find `{"visual", "image_tower", "vision_tower", "vision_model"}` inside both `JanusModel` and `JanusForConditionalGeneration`. For the former it works OK, but it throws errors for the latter, because it should look for `obj.model.vision_model` in this case, not `obj.vision_model`. You can verify that the test fails by commenting the code from `test_sdpa_can_dispatch_composite_models` and calling `super().test_sdpa_can_dispatch_composite_models` instead.
created_at: 2025-03-08T10:41:20Z
author_association: NONE
diff_hunk: @@ -1510,8 +1519,12 @@ def generate(
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None) 
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(input_tokens[batch_size:, :] != generation_config.bos_token_id,
+                                                  generation_config.pad_token_id)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I missed that they were also leaving `boi` because of the -1, will add the changes (ps: the code works even without it). It is necessary to look for the positions instead of just assuming that the first position has `bos` bacause the first will not have `bos` in cases of left padding.
created_at: 2025-03-08T10:45:08Z
author_association: NONE
diff_hunk: @@ -1544,18 +1557,29 @@ def generate(
         decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
+
         for i in tqdm(range(num_image_tokens)):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
-            outputs = self.model.language_model(
+            model_inputs = super().prepare_inputs_for_generation(
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: Will look at what it does more closely and make the change, but wasn't the reason for the new `prepare_inputs_for_generation` something related to the image tokens for the cases where the user is passing an image (which I think it is not in this case)? Can you clarify why was it necessary to override?
created_at: 2025-03-08T10:47:22Z
author_association: NONE
diff_hunk: @@ -1572,6 +1596,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = torch.cat([input_tokens, next_token[:, None]], dim=-1)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I added when I was testing my changes simply to be consistent with how the `GenerationMixin._sample` was doing it, and they do pass it to their `prepare_inputs_for_generation` function. Will double-check if it is indeed unnecessary in this function and remove it accordingly.
created_at: 2025-03-08T10:50:17Z
author_association: NONE
diff_hunk: @@ -1510,8 +1519,12 @@ def generate(
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None) 
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(input_tokens[batch_size:, :] != generation_config.bos_token_id,
+                                                  generation_config.pad_token_id)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: I am really not sure, in ideal case scenario we should not mask bos and boi token according to their codebase but  how is it working even after masking boi for uncodition logits 😅 

> It is necessary to look for the positions instead of just assuming that the first position has bos bacause the first will not have bos in cases of left padding.

Just try to mask everything except boi and bos. If it's working fine then good else need to deep dive and find the optimal logic. This multi prompt case is not defined in their example nor they have anything related to training so we can implement this feature by our logic and just make it work
created_at: 2025-03-08T10:51:39Z
author_association: OWNER
diff_hunk: @@ -1544,18 +1557,29 @@ def generate(
         decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
+
         for i in tqdm(range(num_image_tokens)):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
-            outputs = self.model.language_model(
+            model_inputs = super().prepare_inputs_for_generation(
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: For VLMs we have generally overwrite prepare_input_for_generation function to handle pixel values as it is not needed after first token is generated. I have made the change and pushed it. see if the function works for you and if not then we need to refactor there.
created_at: 2025-03-08T10:54:32Z
author_association: OWNER
diff_hunk: @@ -1544,18 +1557,29 @@ def generate(
         decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
         decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
 
+
         for i in tqdm(range(num_image_tokens)):
-            # Fix me: What to do with attention mask when expanding and repeating input ids.
-            # Should we also modify the attention mask if passed?
-            outputs = self.model.language_model(
+            model_inputs = super().prepare_inputs_for_generation(
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: For VLMs we have generally overwrite prepare_input_for_generation function to handle pixel values as it is not needed after first token is generated. I have made the change and pushed it. see if the function works for you and if not then we need to refactor there.
created_at: 2025-03-08T10:54:35Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: Did not understand what you are proposing. How would we end with `model....model....`?
created_at: 2025-03-08T10:56:26Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: model(janusmodel).vision_model(janus vison model).model(janus vision transformer0... from janusforconditionalgeneration perspective
created_at: 2025-03-08T10:57:56Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: Just looked at siglip2. Their Siglip2Model is instantiated exactly like SiglipModel. WIll paste an excerpt below. This is the implementation we have been instructed to not use:
```python
        # First, initialize the text and vision models with proper attention implementation
        text_model = Siglip2TextModel._from_config(text_config)
        vision_model = Siglip2VisionModel._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
```
created_at: 2025-03-08T10:59:20Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: ouch missed this. Okay let's leave the sdpa case for now as we can always overwrite it based on our final model structure. 
created_at: 2025-03-08T11:01:18Z
author_association: OWNER
diff_hunk: @@ -1510,8 +1519,12 @@ def generate(
         num_image_tokens = self.model.vision_model.config.num_image_tokens
 
         input_tokens = input_ids.repeat(2, 1)  # Double batch size for conditional/unconditional logits
+        attention_mask = model_kwargs.pop("attention_mask", None) 
+        attention_mask = attention_mask.repeat(2, 1)
+
+        input_tokens[batch_size:, :].masked_fill_(input_tokens[batch_size:, :] != generation_config.bos_token_id,
+                                                  generation_config.pad_token_id)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I will make the change, I masked boi because I did not see it. Multi-prompt is not in their examples, but, from my inspections, HF designed things so that, by the use of position ids, the model will see the input like it would if a single image was passed. You can test that the version in this PR works.
created_at: 2025-03-08T11:06:09Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I haven't seen examples with `vision_model.model`. I can keep the name duplicated, as perhaps there is resistance to using `transformer` in the name. I simply used it because of the class name. Should I keep `vision_model.vision_model`? I might be able to avoid overwriting the forward pass like this.
created_at: 2025-03-08T11:10:17Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: Once take a pull from my latest changes, I have changed the attention interface itself and the `test_dispatch_sdpa-composite` testcase is failing 😅 . Haven't deep dived into it to resolve but this one testcase is bugging a lot. So, first lets build the model structure and later override it.

I am a bit hesitant with hacky way we are using rn like del vision model is not right and even the fix using tmp variable is not optimal as in siglip it was used for BC and I don't see any newly merged models have this type of flow.

you can go through some vlm codebases (llava,idefics,paligemma,qwen etc) and try design the vision module similar to how its defined there. WDYT?
created_at: 2025-03-08T11:19:11Z
author_association: OWNER
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: PS: deleting is just a way to avoid the unraveling from including the attribute, they instruct us to do it like this in that modular-module tutorial. The `del ...` line does not appear in the modeling file. And having JanusVisionModel as a proper wrapper saves us from having to delete the tmp like they do.
created_at: 2025-03-08T11:25:33Z
author_association: NONE
diff_hunk: @@ -731,11 +732,13 @@ class JanusVisionModel(SiglipVisionModel):
 
     def __init__(self, config: JanusVisionConfig):
         super().__init__(config)
-
-        self.vision_model = JanusVisionTransformer(config)
-
+        del self.vision_model
+        self.vision_transformer = JanusVisionTransformer(config)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: I will look
created_at: 2025-03-08T11:26:26Z
author_association: NONE
diff_hunk: @@ -1572,6 +1596,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = torch.cat([input_tokens, next_token[:, None]], dim=-1)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: Am glancing at the code now. As I have said `GenerationMixin._sample`---which we use for language model---updates input_ids like this and also, the function `GenerationMixin.prepare_inputs_for_generation`, which we use, even in your new branch, is also making heavy use of input ids. I am still checking there if it is OK to simply ignore this or pass an outdated input_ids every iteration. Passing None throws errors.
created_at: 2025-03-08T13:19:27Z
author_association: NONE
diff_hunk: @@ -1572,6 +1596,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = torch.cat([input_tokens, next_token[:, None]], dim=-1)
path: src/transformers/models/janus/modular_janus.py
user: login: hsilva664
body: It appears this can be changed to `input_tokens = next_token[:, None]`
created_at: 2025-03-08T13:25:29Z
author_association: NONE
diff_hunk: @@ -1572,6 +1596,9 @@ def generate(
             img_embeds = self.prepare_embeddings_for_image_generation(next_token)
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
+            # similar to GenerationMixin._sample, this is needed in prepare_inputs_for_generation
+            input_tokens = torch.cat([input_tokens, next_token[:, None]], dim=-1)
path: src/transformers/models/janus/modular_janus.py
user: login: yaswanth19
body: > input_tokens = next_token[:, None]

Yup, thats what I was saying, we only need to pass new ids internally it should utilize cache for previous context.
created_at: 2025-03-08T13:39:31Z
author_association: OWNER

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-07T23:17:26Z
body: Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page).
author_association: NONE