Update app.py

app.py

@@ -1,18 +1,53 @@
 import gradio as gr
+from PIL import Image
+import spaces
 from transformers import pipeline
 
-get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base")
+# Initialize Model
+get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)
 
-def captioner(input):
+# def captioner(input):
+#     output = get_completion(input)
+#     return output[0]['generated_text']
+
+@spaces.GPU(duration=120)
+def captioner(input: Image.Image) -> str:
+    """
+    Generate a caption for the given image using the BLIP-IMAGE-CAPTIONING-BASE model.
+
+    Args:
+        input (Image.Image): The input image for which to generate a caption.
+
+    Returns:
+        str: The generated caption text.
+    """
     output = get_completion(input)
     return output[0]['generated_text']
 
 gr.close_all()
-demo = gr.Interface(fn=captioner,
+
+####### GRADIO APP #######
+title = """<h1 id="title"> Image Captioning with BLIP model </h1>"""
+
+description = """
+- The model used for Generating Captions [BLIP-IMAGE-CAPTIONING-BASE](https://huggingface.co/Salesforce/blip-image-captioning-base).
+"""
+
+css = '''
+h1#title {
+  text-align: center;
+}
+'''
+
+theme = gr.themes.Soft()
+demo = gr.Blocks(css=css, theme=theme)
+
+with demo:
+  gr.Markdown(title)
+  gr.Markdown(description)
+  interface = gr.Interface(fn=captioner,
                     inputs=[gr.Image(label="Upload image", type="pil")],
                     outputs=[gr.Textbox(label="Caption")],
-                    title="Image Captioning with BLIP model",
-                    description="Caption any image using the `Salesforce/blip-image-captioning-base` model ",
                     allow_flagging="never")
 
 demo.launch()