Update app.py

app.py

@@ -1,4 +1,3 @@
-# type: ignore
 from typing import Any
 import gradio as gr
 import spaces
@@ -6,7 +5,6 @@ import torch
 from PIL import Image
 from transformers import AutoModelForCausalLM, LlamaTokenizer
 
-
 DEFAULT_PARAMS = {
     "do_sample": False,
     "max_new_tokens": 256,
@@ -22,7 +20,7 @@ DEFAULT_QUERY = (
     "Avoid subjective interpretations or speculation."
 )
 
-DTYPE = torch.bfloat16
+DTYPE = torch.float16  # Use float16 for faster processing on CPU with limited resources
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = LlamaTokenizer.from_pretrained(
@@ -58,25 +56,23 @@ def generate_caption(
     }
 
     outputs = model.generate(**inputs, **params)
-    outputs = outputs[:, inputs["input_ids"].shape[1] :]
+    outputs = outputs[:, inputs["input_ids"].shape[1]:]
     result = tokenizer.decode(outputs[0])
 
     result = result.replace("This image showcases", "").strip().removesuffix("</s>").strip().capitalize()
     return result
 
-
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil")
-            input_query = gr.Textbox(lines=5, label="Prompt", value=DEFAULT_QUERY)
             run_button = gr.Button(value="Generate Caption")
         with gr.Column():
             output_caption = gr.Textbox(label="Generated Caption", show_copy_button=True)
 
     run_button.click(
         fn=generate_caption,
-        inputs=[input_image, input_query],
+        inputs=[input_image],  # Only input image is needed
         outputs=output_caption,
     )