first commit

app.py

@@ -6,90 +6,75 @@ import os
 
 hf_token = os.environ["HF_TOKEN"]
 
+# Load the Gemma 3 pipeline
+pipe = pipeline(
+    "image-text-to-text",
+    model="google/gemma-3-4b-it",
+    device="cuda",
+    torch_dtype=torch.bfloat16,
+    use_auth_token=hf_token
+)
+
 @spaces.GPU
-def get_response(message, chat_history, image=None):
-    # Choose the appropriate pipeline based on whether an image is provided
-    if image is not None:
-        # Multimodal pipeline for text+image
-        pipe = pipeline(
-            "image-text-to-text",
-            model="google/gemma-3-4b-it",
-            device="cuda",
-            torch_dtype=torch.bfloat16,
-            use_auth_token=hf_token
-        )
-        
-        messages = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": "You are a helpful assistant."}]
-            }
-        ]
-        
-        user_content = []
-        user_content.append({"type": "image", "image": image})
-        if message:
-            user_content.append({"type": "text", "text": message})
-            
-        messages.append({"role": "user", "content": user_content})
-        
-    else:
-        # Text-only pipeline
-        pipe = pipeline(
-            "text-generation",
-            model="google/gemma-3-4b-it",
-            device="cuda",
-            torch_dtype=torch.bfloat16,
-            use_auth_token=hf_token
-        )
+def get_response(message, chat_history, image):
+    # Check if image is provided
+    if image is None:
+        chat_history.append((message, "Please upload an image (required)"))
+        return "", chat_history
+    
+    messages = [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "You are a helpful assistant."}]
+        }
+    ]
+    
+    user_content = [{"type": "image", "image": image}]
+    
+    # Add text message if provided
+    if message:
+        user_content.append({"type": "text", "text": message})
         
-        messages = [
-            {
-                "role": "system",
-                "content": "You are a helpful assistant."
-            },
-            {
-                "role": "user",
-                "content": message
-            }
-        ]
+    messages.append({"role": "user", "content": user_content})
     
-    # Call the appropriate pipeline
+    # Call the pipeline
     output = pipe(text=messages, max_new_tokens=200)
     
     try:
-        if image is not None:
-            response = output[0]["generated_text"][-1]["content"]
-        else:
-            response = output[0]["generated_text"]
-        
+        response = output[0]["generated_text"][-1]["content"]
         chat_history.append((message, response))
-        return "", chat_history
     except (KeyError, IndexError, TypeError) as e:
         error_message = f"Error processing the response: {str(e)}"
         chat_history.append((message, error_message))
-        return "", chat_history
+    
+    return "", chat_history
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Gemma 3 Chat Interface")
-    gr.Markdown("Chat with Gemma 3 with optional image upload capability")
+    gr.Markdown("# Gemma 3 Image Chat")
+    gr.Markdown("Chat with Gemma 3 about images. Image upload is required for each message.")
     
     chatbot = gr.Chatbot()
     
     with gr.Row():
         msg = gr.Textbox(
             show_label=False,
-            placeholder="Type your message here...",
+            placeholder="Type your message here about the image...",
             scale=4
         )
         img = gr.Image(
             type="pil", 
-            label="Upload image (optional)", 
+            label="Upload image (required)", 
             scale=1
         )
     
     submit_btn = gr.Button("Send")
     
+    # Clear button to reset the interface
+    clear_btn = gr.Button("Clear")
+    
+    def clear_interface():
+        return "", [], None
+    
     submit_btn.click(
         get_response,
         inputs=[msg, chatbot, img],
@@ -101,6 +86,12 @@ with gr.Blocks() as demo:
         inputs=[msg, chatbot, img],
         outputs=[msg, chatbot]
     )
+    
+    clear_btn.click(
+        clear_interface,
+        inputs=None,
+        outputs=[msg, chatbot, img]
+    )
 
 if __name__ == "__main__":
     demo.launch()