Update app.py

app.py

@@ -5,6 +5,7 @@ import torch
 from PIL import Image
 from transformers import AutoModelForCausalLM, LlamaTokenizer
 
+# Constants
 DEFAULT_PARAMS = {
     "do_sample": False,
     "max_new_tokens": 256,
@@ -23,6 +24,7 @@ DEFAULT_QUERY = (
 DTYPE = torch.bfloat16
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Load model and tokenizer
 tokenizer = LlamaTokenizer.from_pretrained(
     pretrained_model_name_or_path="lmsys/vicuna-7b-v1.5",
 )
@@ -43,7 +45,7 @@ def generate_caption(
 ) -> str:
     inputs = model.build_conversation_input_ids(
         tokenizer=tokenizer,
-        query=DEFAULT_QUERY,  # Use the default query directly
+        query=DEFAULT_QUERY,
         history=[],
         images=[image],
     )
@@ -61,13 +63,41 @@ def generate_caption(
     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")  # Image input remains
-            run_button = gr.Button(value="Generate Caption")
-        with gr.Column():
-            output_caption = gr.Textbox(label="Generated Caption", show_copy_button=True)
+# CSS for design enhancements with a dark button and white text
+css = """
+  #container {
+    background-color: #f9f9f9;
+    padding: 20px;
+    border-radius: 15px;
+    border: 2px solid #333; /* Darker outline */
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2); /* Enhanced shadow */
+    max-width: 400px;
+    margin: auto;
+  }
+  #input_image, #output_caption, #run_button {
+    margin-top: 15px;
+    border: 2px solid #333; /* Darker outline */
+    border-radius: 8px;
+  }
+  #run_button {
+    background-color: #000000; /* Dark button color */
+    color: white; /* White text */
+    border-radius: 10px;
+    padding: 10px;
+    cursor: pointer;
+    transition: background-color 0.3s ease;
+  }
+  #run_button:hover {
+    background-color: #333; /* Slightly lighter on hover */
+  }
+"""
+
+# Gradio interface with vertical alignment
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="container"):
+        input_image = gr.Image(type="pil", elem_id="input_image")
+        run_button = gr.Button(value="Generate Caption", elem_id="run_button")
+        output_caption = gr.Textbox(label="Generated Caption", show_copy_button=True, elem_id="output_caption")
 
     run_button.click(
         fn=generate_caption,