Update app.py

app.py

@@ -11,43 +11,29 @@ model, _, transform = open_clip.create_model_and_transforms(
     pretrained="laion2b_s13b_b90k"
 )
 model.to(device)
-model.eval()  # Ensure the model is in evaluation mode
+model.eval()
 
-def output_generate(image_batch):
+def output_generate(image):
+    im = transform(image).unsqueeze(0).to(device)
     with torch.no_grad(), torch.cuda.amp.autocast():
-        generated = model.generate(image_batch.to(device), seq_len=20)
-    captions = open_clip.decode(generated.detach()).split("<end_of_text>")
-    captions = [c.replace("<start_of_text>", "") for c in captions]
-    return captions
+        generated = model.generate(im, seq_len=20)
+    return open_clip.decode(generated[0].detach()).split("<end_of_text>")[0].replace("<start_of_text>", "")
 
-def inference_caption(image_batch):
+def inference_caption(image):
+    im = transform(image).unsqueeze(0).to(device)
     with torch.no_grad(), torch.cuda.amp.autocast():
         generated = model.generate(
-            image_batch.to(device), 
+            im, 
             generation_type="beam_search",
             top_p=1.0, 
             min_seq_len=20, 
             seq_len=30, 
             repetition_penalty=1.2
         )
-    captions = open_clip.decode(generated.detach()).split("<end_of_text>")
-    captions = [c.replace("<start_of_text>", "") for c in captions]
-    return captions
+    return open_clip.decode(generated[0].detach()).split("<end_of_text>")[0].replace("<start_of_text>", "")
 
-# Use Gradio's batching feature to process multiple images at once
-image_input = gr.inputs.Image(type="pil", label="Input Image", max_shape=(224, 224))
+image_input = gr.inputs.Image(type="pil")
 caption_output = gr.outputs.Textbox(label="Caption Output")
-caption_interface = gr.Interface(
-    fn=inference_caption, 
-    inputs=gr.inputs.Image(type="pil", label="Input Image", max_shape=(224, 224), multiple=True), 
-    outputs=gr.outputs.Textbox(label="Caption Output", type="auto", multiple=True),
-    capture_session=True, 
-    title="CoCa: Contrastive Captioners", 
-    description="An open source implementation of CoCa: Contrastive Captioners are Image-Text Foundation Models.", 
-    examples=[path.as_posix() for path in sorted(pathlib.Path("images").glob("*.jpg"))], 
-    allow_flagging=False,
-    batching=True,  # Enable Gradio's batching feature
-    batch_size=8,  # Process 8 images at once
-)
+caption_interface = gr.Interface(fn=inference_caption, inputs=image_input, outputs=caption_output, capture_session=True, title="CoCa: Contrastive Captioners", description="An open source implementation of CoCa: Contrastive Captioners are Image-Text Foundation Models.", examples=[path.as_posix() for path in sorted(pathlib.Path("images").glob("*.jpg"))], allow_flagging=False)
 
-caption_interface.launch()
+caption_interface.launch()