duplicaction test

app.py

@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 import gradio as gr
 import torch
 from transformers import AutoProcessor, LlavaForConditionalGeneration
@@ -24,10 +26,16 @@ model = LlavaForConditionalGeneration.from_pretrained(
 )
 
 
-def text_to_image(image, prompt):
+def text_to_image(image, prompt, duplications: int):
     prompt = f'USER: <image>\n{prompt}\nASSISTANT:'
 
-    inputs = processor([prompt], images=[image], padding=True, return_tensors="pt").to(model.device)
+    image_batch = [image]
+    prompt_batch = [prompt]
+    for _ in range(duplications):
+        image_batch.append(deepcopy(image))
+        prompt_batch.append(prompt)
+
+    inputs = processor(prompt_batch, images=image_batch, padding=True, return_tensors="pt").to(model.device)
     output = model.generate(**inputs, max_new_tokens=500, temperature=0.3)
     generated_text = processor.batch_decode(output, skip_special_tokens=True)
     text = generated_text.pop()
@@ -41,7 +49,8 @@ demo = gr.Interface(
     fn=text_to_image,
     inputs=[
         gr.Image(label='Select an image to analyze', type='pil'),
-        gr.Textbox(label='Enter Prompt')
+        gr.Textbox(label='Enter Prompt'),
+        gr.Number(label='How many duplications of the image (to test memory load)', value=0)
     ],
     outputs=[gr.Textbox(label='Maurice says:'), gr.JSON(label='Embedded text')]
 )