Update cumo/model/language_model/llava_llama.py

cumo/model/language_model/llava_llama.py

@@ -107,7 +107,7 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
     @torch.no_grad()
     def generate(
         self,
-        inputs: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor] = None,
         images: Optional[torch.Tensor] = None,
         image_sizes: Optional[torch.Tensor] = None,
         **kwargs,
@@ -116,30 +116,28 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
         attention_mask = kwargs.pop("attention_mask", None)
         if "inputs_embeds" in kwargs:
             raise NotImplementedError("`inputs_embeds` is not supported")
-
-        if inputs_embeds is None:
+    
+        if images is not None:
             (
                 input_ids,
                 position_ids,
                 attention_mask,
-                past_key_values,
+                _,
                 inputs_embeds,
-                labels,
+                _,
                 *_
             ) = self.prepare_inputs_labels_for_multimodal(
                 input_ids,
                 position_ids,
                 attention_mask,
-                past_key_values,
-                labels,
+                None,
+                None,
                 images,
-                image_sizes
+                image_sizes=image_sizes
             )
-
-
         else:
-            inputs_embeds = self.get_model().embed_tokens(inputs)
-
+            inputs_embeds = self.get_model().embed_tokens(input_ids)
+    
         return super().generate(
             position_ids=position_ids,
             attention_mask=attention_mask,
@@ -147,6 +145,7 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
             **kwargs
         )
 
+
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                       inputs_embeds=None, **kwargs):
         images = kwargs.pop("images", None)