first commit

app.py

@@ -2,15 +2,13 @@ import spaces
 import gradio as gr
 import torch
 from transformers import AutoProcessor, AutoModelForImageTextToText
-from PIL import Image
-import io
 import requests
 import os
 
 hf_token = os.environ.get("HF_TOKEN")
 model_id = "CohereForAI/aya-vision-8b"
 
-# Load model and processor immediately on startup.
+# Load the model and processor during startup.
 try:
     processor = AutoProcessor.from_pretrained(model_id)
     model = AutoModelForImageTextToText.from_pretrained(
@@ -26,36 +24,34 @@ except Exception as e:
     )
 
 @spaces.GPU
-def process_image_and_prompt(image, image_url, prompt, temperature=0.3, max_tokens=300):
+def process_image_and_prompt(uploaded_image, image_url, prompt, temperature=0.3, max_tokens=300):
     global processor, model
 
-    # Check if the model is loaded
     if processor is None or model is None:
         return "Model failed to load. Please check the logs."
-
-    # Process image input (either uploaded or from URL)
-    if image is not None:
-        img = Image.fromarray(image)
+    
+    # Determine which image to use:
+    # If an image is uploaded, it is returned as a file path.
+    if uploaded_image is not None:
+        # If the file path does not start with "http", prefix with '/file/' so that
+        # the Hugging Face Space can serve it via an HTTP URL.
+        img_url = uploaded_image if uploaded_image.startswith("http") else f"/file/{uploaded_image}"
     elif image_url and image_url.strip():
-        try:
-            response = requests.get(image_url)
-            img = Image.open(io.BytesIO(response.content))
-        except Exception as e:
-            return f"Error loading image from URL: {e}"
+        img_url = image_url.strip()
     else:
-        return "Please provide either an image or an image URL."
-
-    # Format the message using the Aya Vision chat template
+        return "Please provide either an image upload or an image URL."
+    
+    # Build the message using the Aya Vision chat template.
     messages = [
         {
             "role": "user",
             "content": [
-                {"type": "image", "source": img},
+                {"type": "image", "url": img_url},
                 {"type": "text", "text": prompt},
             ],
         },
     ]
-
+    
     try:
         inputs = processor.apply_chat_template(
             messages,
@@ -65,29 +61,30 @@ def process_image_and_prompt(image, image_url, prompt, temperature=0.3, max_toke
             return_dict=True,
             return_tensors="pt"
         ).to(model.device)
-
+    
         gen_tokens = model.generate(
             **inputs,
             max_new_tokens=int(max_tokens),
             do_sample=True,
             temperature=float(temperature),
         )
-
+    
         response = processor.tokenizer.decode(
-            gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True
+            gen_tokens[0][inputs.input_ids.shape[1]:],
+            skip_special_tokens=True
         )
         return response
     except Exception as e:
         return f"Error generating response: {e}"
 
-# Example inputs for testing
+# Example inputs for testing.
 examples = [
     [None, "https://media.istockphoto.com/id/458012057/photo/istanbul-turkey.jpg?s=612x612&w=0&k=20&c=qogAOVvkpfUyqLUMr_XJQyq-HkACXyYUSZbKhBlPrxo=", "What landmark is shown in this image?", 0.3, 300],
     [None, "https://pbs.twimg.com/media/Fx7YvfQWYAIp6rZ?format=jpg&name=medium", "What does the text in this image say?", 0.3, 300],
     [None, "https://upload.wikimedia.org/wikipedia/commons/d/da/The_Parthenon_in_Athens.jpg", "Describe esta imagen en español", 0.3, 300]
 ]
 
-# Create the Gradio application
+# Build the Gradio interface.
 with gr.Blocks(title="Aya Vision 8B Demo") as demo:
     gr.Markdown("# Aya Vision 8B Model Demo")
     gr.Markdown("""
@@ -101,46 +98,44 @@ with gr.Blocks(title="Aya Vision 8B Demo") as demo:
     Upload an image or provide a URL, and enter a prompt to get started!
     """)
     
-    # Display model loading status
+    # Display model loading status.
     gr.Markdown(f"**Model Status:** {model_status}")
-
-    gr.Markdown("### Upload an image or provide an image URL:")
+    
+    gr.Markdown("### Provide an image (upload or URL):")
     with gr.Tab("Upload Image"):
-        image_input = gr.Image(label="Upload Image", type="numpy")
-        image_url_input = gr.Textbox(label="Image URL", placeholder="Leave blank if uploading an image", visible=False)
+        # Set type to 'filepath' to get the file path from the upload.
+        image_input = gr.Image(label="Upload Image", type="filepath")
     with gr.Tab("Image URL"):
-        image_url_visible = gr.Textbox(label="Image URL", placeholder="Enter a URL to an image")
-        image_input_url = gr.Image(label="Upload Image", type="numpy", visible=False)
-
+        image_url_input = gr.Textbox(label="Image URL", placeholder="Enter a URL to an image")
+    
     prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt to the model", lines=3)
-
+    
     with gr.Accordion("Generation Settings", open=False):
         temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.3, label="Temperature")
         max_tokens = gr.Slider(minimum=50, maximum=1000, step=50, value=300, label="Max Tokens")
-
+    
     generate_button = gr.Button("Generate Response", variant="primary")
-
+    
     with gr.Column():
         output = gr.Textbox(label="Model Response", lines=10)
-
+    
     gr.Markdown("### Examples")
     gr.Examples(
         examples=examples,
-        inputs=[image_input, image_url_visible, prompt, temperature, max_tokens],
+        inputs=[image_input, image_url_input, prompt, temperature, max_tokens],
         outputs=output,
         fn=process_image_and_prompt
     )
-
-    # Handle generation from either image or URL
-    def generate_response(image, image_url_visible, prompt, temperature, max_tokens):
-        return process_image_and_prompt(image, image_url_visible, prompt, temperature, max_tokens)
-
+    
+    # Determine which image input to use when generating the response.
+    def generate_response(uploaded_image, image_url, prompt, temperature, max_tokens):
+        return process_image_and_prompt(uploaded_image, image_url, prompt, temperature, max_tokens)
+    
     generate_button.click(
         generate_response,
-        inputs=[image_input, image_url_visible, prompt, temperature, max_tokens],
+        inputs=[image_input, image_url_input, prompt, temperature, max_tokens],
         outputs=output
     )
 
-# Launch the Gradio app
 if __name__ == "__main__":
     demo.launch()