first commit

app.py

@@ -6,47 +6,69 @@ import os
 
 hf_token = os.environ["HF_TOKEN"]
 
-# Load the Gemma 3 pipeline - use the multimodal version for all cases
-pipe = pipeline(
-    "image-text-to-text",  # This pipeline can handle both text-only and text+image
-    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):
-    messages = [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": "You are a helpful assistant."}]
-        }
-    ]
-    
-    user_content = []
-    
-    # Only add image if provided
+    # 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})
-    
-    # Always add the text message
-    if message:
-        user_content.append({"type": "text", "text": message})
+        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
+        )
         
-    messages.append({"role": "user", "content": user_content})
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant."
+            },
+            {
+                "role": "user",
+                "content": message
+            }
+        ]
     
-    # Call the pipeline
+    # Call the appropriate pipeline
     output = pipe(text=messages, max_new_tokens=200)
     
     try:
-        response = output[0]["generated_text"][-1]["content"]
+        if image is not None:
+            response = output[0]["generated_text"][-1]["content"]
+        else:
+            response = output[0]["generated_text"]
+        
         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")