Merge pull request #63 from rhymes-ai/pixel_mask

fix: add pixel_mask support to generation

aria/model/modeling_aria.py

@@ -332,6 +332,7 @@ def prepare_inputs_for_generation(
         past_key_values=None,
         inputs_embeds=None,
         pixel_values=None,
+        pixel_mask=None,
         attention_mask=None,
         cache_position=None,
         num_logits_to_keep=None,
@@ -351,5 +352,6 @@ def prepare_inputs_for_generation(
             # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
             # Otherwise we need pixel values to be passed to model
             model_inputs["pixel_values"] = pixel_values
+            model_inputs["pixel_mask"] = pixel_mask
 
         return model_inputs