Update app.py

app.py

@@ -20,7 +20,7 @@ DEFAULT_QUERY = (
     "Avoid subjective interpretations or speculation."
 )
 
-DTYPE = torch.float16  # Use float16 for faster processing on CPU with limited resources
+DTYPE = torch.bfloat16
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = LlamaTokenizer.from_pretrained(
@@ -39,12 +39,11 @@ model = model.to(device=DEVICE)
 @torch.no_grad()
 def generate_caption(
     image: Image.Image,
-    query: str = DEFAULT_QUERY,
     params: dict[str, Any] = DEFAULT_PARAMS,
 ) -> str:
     inputs = model.build_conversation_input_ids(
         tokenizer=tokenizer,
-        query=query,
+        query=DEFAULT_QUERY,  # Use the default query directly
         history=[],
         images=[image],
     )
@@ -56,7 +55,7 @@ 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()
@@ -65,14 +64,14 @@ def generate_caption(
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="pil")
+            input_image = gr.Image(type="pil")  # Image input remains
             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],  # Only input image is needed
+        inputs=[input_image],  # Only the image input is passed
         outputs=output_caption,
     )