Update app.py

app.py

@@ -5,12 +5,13 @@ import torch
 from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
 
 # Model and Processor Setup
-model_id = "google/paligemma2-3b-mix-448"
+model_id = "gv-hf/paligemma2-3b-mix-448"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 HF_KEY = os.getenv("HF_KEY")
 if not HF_KEY:
     raise ValueError("Please set the HF_KEY environment variable with your Hugging Face API token")
 
+# Load model and processor
 model = PaliGemmaForConditionalGeneration.from_pretrained(
     model_id,
     token=HF_KEY,
@@ -47,25 +48,8 @@ def detect_objects(image: PIL.Image.Image) -> str:
 def vqa(image: PIL.Image.Image, question: str) -> str:
     return infer(image, f"Q: {question} A:", max_new_tokens=50)
 
-# Custom CSS for Styling
-custom_css = """
-.gradio-container {
-    font-family: 'Arial', sans-serif;
-}
-.upload-button {
-    background-color: #4285f4;
-    color: white;
-    border-radius: 5px;
-    padding: 10px 20px;
-}
-.output-text {
-    font-size: 18px;
-    font-weight: bold;
-}
-"""
-
 # Gradio App
-with gr.Blocks(css=custom_css) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# PaliGemma Multi-Modal App")
     gr.Markdown("Upload an image and explore its features using the PaliGemma model!")
 
@@ -75,9 +59,9 @@ with gr.Blocks(css=custom_css) as demo:
             with gr.Row():
                 with gr.Column():
                     caption_image = gr.Image(type="pil", label="Upload Image", width=512, height=512)
-                    caption_btn = gr.Button("Generate Caption", elem_classes="upload-button")
+                    caption_btn = gr.Button("Generate Caption")
                 with gr.Column():
-                    caption_output = gr.Text(label="Generated Caption", elem_classes="output-text")
+                    caption_output = gr.Text(label="Generated Caption")
             caption_btn.click(fn=generate_caption, inputs=[caption_image], outputs=[caption_output])
 
         # Tab 2: Object Detection
@@ -85,9 +69,9 @@ with gr.Blocks(css=custom_css) as demo:
             with gr.Row():
                 with gr.Column():
                     detect_image = gr.Image(type="pil", label="Upload Image", width=512, height=512)
-                    detect_btn = gr.Button("Detect Objects", elem_classes="upload-button")
+                    detect_btn = gr.Button("Detect Objects")
                 with gr.Column():
-                    detect_output = gr.Text(label="Detected Objects", elem_classes="output-text")
+                    detect_output = gr.Text(label="Detected Objects")
             detect_btn.click(fn=detect_objects, inputs=[detect_image], outputs=[detect_output])
 
         # Tab 3: Visual Question Answering (VQA)
@@ -96,9 +80,9 @@ with gr.Blocks(css=custom_css) as demo:
                 with gr.Column():
                     vqa_image = gr.Image(type="pil", label="Upload Image", width=512, height=512)
                     vqa_question = gr.Text(label="Ask a Question", placeholder="What is in the image?")
-                    vqa_btn = gr.Button("Ask", elem_classes="upload-button")
+                    vqa_btn = gr.Button("Ask")
                 with gr.Column():
-                    vqa_output = gr.Text(label="Answer", elem_classes="output-text")
+                    vqa_output = gr.Text(label="Answer")
             vqa_btn.click(fn=vqa, inputs=[vqa_image, vqa_question], outputs=[vqa_output])
 
         # Tab 4: Text Generation (Original Feature)
@@ -107,21 +91,11 @@ with gr.Blocks(css=custom_css) as demo:
                 with gr.Column():
                     text_image = gr.Image(type="pil", label="Upload Image", width=512, height=512)
                     text_input = gr.Text(label="Input Text", placeholder="Describe the image...")
-                    text_btn = gr.Button("Generate Text", elem_classes="upload-button")
+                    text_btn = gr.Button("Generate Text")
                 with gr.Column():
-                    text_output = gr.Text(label="Generated Text", elem_classes="output-text")
+                    text_output = gr.Text(label="Generated Text")
             text_btn.click(fn=infer, inputs=[text_image, text_input, gr.Slider(10, 200, value=50)], outputs=[text_output])
 
-    # Image Upload/Download
-    with gr.Row():
-        upload_button = gr.UploadButton("Upload Image", file_types=["image"], elem_classes="upload-button")
-        download_button = gr.DownloadButton("Download Results", elem_classes="upload-button")
-
-    # Real-Time Updates
-    caption_image.change(fn=generate_caption, inputs=[caption_image], outputs=[caption_output], live=True)
-    detect_image.change(fn=detect_objects, inputs=[detect_image], outputs=[detect_output], live=True)
-    vqa_image.change(fn=lambda x: vqa(x, "What is in the image?"), inputs=[vqa_image], outputs=[vqa_output], live=True)
-
 # Launch the App
 if __name__ == "__main__":
     demo.queue(max_size=10).launch(debug=True)