Update spaces to use ZeroGPU

app.py

@@ -1,9 +1,9 @@
 import os
-from dotenv import load_dotenv
 import gradio as gr
-import torch
-from PIL import Image
 from transformers import pipeline
+import torch
+from dotenv import load_dotenv
+import spaces
 
 load_dotenv()
 hf_token = os.getenv("HF_GEMMA_TOKEN")
@@ -11,16 +11,18 @@ hf_token = os.getenv("HF_GEMMA_TOKEN")
 pipe = pipeline(
     "image-text-to-text",
     model="google/gemma-3-4b-it",
-    device_map="cpu",
-    torch_dtype=torch.float32,
-    token=hf_token
+    token=hf_token,
+    device="cuda",
+    torch_dtype=torch.bfloat16,
 )
 
-def analyze_image(image):
-    image = image.convert('RGB')
+@spaces.GPU()
+def get_response(chat_history, image):
+    if image is None:
+        chat_history.append(("Please upload an image (required)", ""))
+        return chat_history
     
     # Define the prompt
-    system_prompt = "You are a helpful assistant."
     user_prompt = '''Analyze this food image and provide detailed nutritional information in JSON format. 
     Identify the specific vegetarian food items, estimate portion sizes, and provide nutritional breakdown.
     This app focuses on vegetarian foods only, so analyze from that perspective.
@@ -45,11 +47,11 @@ def analyze_image(image):
     }
     
     Return ONLY the JSON without any explanations or markdown formatting.'''
-    
+        
     messages = [
         {
             "role": "system",
-            "content": [{"type": "text", "text": system_prompt}]
+            "content": [{"type": "text", "text": "You are a helpful assistant."}]
         },
         {
             "role": "user",
@@ -60,19 +62,48 @@ def analyze_image(image):
         }
     ]
     
-    output = pipe(text=messages, max_new_tokens=500, return_full_text=False)
+    output = pipe(text=messages, max_new_tokens=200)
+    
+    try:
+        response = output[0]["generated_text"][-1]["content"]
+        chat_history.append((user_prompt, response))
+    except (KeyError, IndexError, TypeError) as e:
+        error_message = f"Error processing the response: {str(e)}"
+        chat_history.append((user_prompt, error_message))
     
-    return output[0]["generated_text"]
+    return chat_history
 
+with gr.Blocks() as demo:
+    gr.Markdown("# Gemma Powered Calorie Tracker")
+    gr.Markdown("Upload an image to get detailed nutritional information.")
 
-# Gradio interface
-demo = gr.Interface(
-    fn=analyze_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(),
-    title="Gemma Powered Calorie Tracker",
-    description="Upload an image of your food to get detailed nutritional information."
-)
+    chatbot = gr.Chatbot()
+    
+    with gr.Row():
+        img = gr.Image(
+            type="pil", 
+            label="Upload image (required)", 
+            scale=1
+        )
+    
+    submit_btn = gr.Button("Send")
+    
+    clear_btn = gr.Button("Clear")
+    
+    def clear_interface():
+        return [], None
+    
+    submit_btn.click(
+        get_response,
+        inputs=[chatbot, img],
+        outputs=chatbot
+    )
+    
+    clear_btn.click(
+        clear_interface,
+        inputs=None,
+        outputs=[chatbot, img]
+    )
 
 if __name__ == "__main__":
     demo.launch() 

requirements.txt

@@ -1,7 +1,9 @@
-transformers
-torch
-Pillow
-gradio
-accelerate
-transformers @ git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3
-python-dotenv
+git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3
+gradio>=4.0.0
+torch>=2.0.0
+torchvision>=0.15.0
+pillow>=9.0.0
+requests>=2.28.0
+numpy>=1.22.0
+python-dotenv
+spaces