.

utils.py

@@ -71,10 +71,11 @@ def aspect_ratio_preserving_resize_and_crop(image, target_width, target_height):
 
 
 class Image2Text:
-    def __init__(self, model_path, hf_token, device):
+    def __init__(self, model_path, hf_token, device, max_length=1024):
         self.device = device
         self.hf_token = hf_token
         self.model_path = model_path
+        self.max_length = max_length
         self.model, self.processor = self.load_model(self.model_path)
         self.decoder_input_ids = torch.tensor([[self.model.config.decoder_start_token_id]]).to(self.device)
 
@@ -97,7 +98,7 @@ class Image2Text:
         outputs = self.model.generate(
                 pixel_values,
                 decoder_input_ids=self.decoder_input_ids.repeat(pixel_values.shape[0], 1),
-                max_length=1024,
+                max_length=self.max_length,
                 early_stopping=True,
                 pad_token_id=self.processor.tokenizer.pad_token_id,
                 eos_token_id=self.processor.tokenizer.eos_token_id,