Workaround for scaling bug in transformers

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-import os
 import torch
 from transformers import FuyuForCausalLM, AutoTokenizer
 from transformers.models.fuyu.processing_fuyu import FuyuProcessor
 from transformers.models.fuyu.image_processing_fuyu import FuyuImageProcessor
+from PIL import Image
 
 model_id = "adept/fuyu-8b"
 revision = "refs/pr/3"
@@ -16,8 +16,21 @@ processor = FuyuProcessor(image_processor=FuyuImageProcessor(), tokenizer=tokeni
 
 caption_prompt = "Generate a coco-style caption.\\n"
 
+def resize_to_max(image, max_width=1920, max_height=1080):
+    width, height = image.size
+    if width <= max_width and height <= max_height:
+        return image
+
+    scale = min(max_width/width, max_height/height)
+    width = int(width*scale)
+    height = int(height*scale)
+
+    return image.resize((width, height), Image.LANCZOS)
+
 def predict(image, prompt):
     # image = image.convert('RGB')
+    image = resize_to_max(image)
+
     model_inputs = processor(text=prompt, images=[image])
     model_inputs = {k: v.to(dtype=dtype if torch.is_floating_point(v) else v.dtype, device=device) for k,v in model_inputs.items()}
 
@@ -57,7 +70,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Tab("Visual Question Answering"):
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(label="Upload your Image")
+                image_input = gr.Image(label="Upload your Image", type="pil")
                 text_input = gr.Textbox(label="Ask a Question")
             vqa_output = gr.Textbox(label="Output")
             
@@ -75,7 +88,7 @@ with gr.Blocks(css=css) as demo:
         
     with gr.Tab("Image Captioning"):
         with gr.Row():
-            captioning_input = gr.Image(label="Upload your Image")
+            captioning_input = gr.Image(label="Upload your Image", type="pil")
             captioning_output = gr.Textbox(label="Output")
         captioning_btn = gr.Button("Generate Caption")