Update cumo/model/language_model/llava_llama.py

cumo/model/language_model/llava_llama.py

@@ -72,24 +72,25 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
 
         if inputs_embeds is None:
             (
-                inputs,
+                input_ids,
                 position_ids,
                 attention_mask,
-                _,
+                past_key_values,
                 inputs_embeds,
-                _,
+                labels,
                 *_
             ) = self.prepare_inputs_labels_for_multimodal(
-                inputs,
+                input_ids,
                 position_ids,
                 attention_mask,
-                None,
-                None,
+                past_key_values,
+                labels,
                 images,
-                image_sizes=image_sizes
+                image_sizes
             )
 
 
+
         return super().forward(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -116,25 +117,26 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
         if "inputs_embeds" in kwargs:
             raise NotImplementedError("`inputs_embeds` is not supported")
 
-        if images is not None:
+        if inputs_embeds is None:
             (
-                inputs,
+                input_ids,
                 position_ids,
                 attention_mask,
-                _,
+                past_key_values,
                 inputs_embeds,
-                _,
+                labels,
                 *_
             ) = self.prepare_inputs_labels_for_multimodal(
-                inputs,
+                input_ids,
                 position_ids,
                 attention_mask,
-                None,
-                None,
+                past_key_values,
+                labels,
                 images,
-                image_sizes=image_sizes
+                image_sizes
             )
 
+
         else:
             inputs_embeds = self.get_model().embed_tokens(inputs)