first commit

app.py

@@ -4,37 +4,35 @@ import torch
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from PIL import Image
 import io
-import requests 
+import requests
 import os
 
 hf_token = os.environ.get("HF_TOKEN")
-# Initialize model and processor globally for caching
 model_id = "CohereForAI/aya-vision-8b"
-processor = None
-model = None
 
-@spaces.GPU
-def load_model():
-    global processor, model
-    if processor is None or model is None:
-        try:
-            processor = AutoProcessor.from_pretrained(model_id)
-            model = AutoModelForImageTextToText.from_pretrained(
-                model_id, device_map="auto", torch_dtype=torch.float16, use_auth_token=hf_token
-            )
-            return "Model loaded successfully!"
-        except Exception as e:
-            return f"Error loading model: {e}\nMake sure to install the correct version of transformers with: pip install 'git+https://github.com/huggingface/[email protected]'"
-    return "Model already loaded!"
+# Load model and processor immediately on startup.
+try:
+    processor = AutoProcessor.from_pretrained(model_id)
+    model = AutoModelForImageTextToText.from_pretrained(
+        model_id, device_map="auto", torch_dtype=torch.float16, use_auth_token=hf_token
+    )
+    model_status = "Model loaded successfully!"
+except Exception as e:
+    processor = None
+    model = None
+    model_status = (
+        f"Error loading model: {e}\nMake sure to install the correct version of transformers with: "
+        "pip install 'git+https://github.com/huggingface/[email protected]'"
+    )
 
 @spaces.GPU
 def process_image_and_prompt(image, image_url, prompt, temperature=0.3, max_tokens=300):
     global processor, model
-    
-    # Ensure model is loaded
+
+    # Check if the model is loaded
     if processor is None or model is None:
-        return "Please load the model first using the 'Load Model' button."
-    
+        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)
@@ -46,48 +44,50 @@ def process_image_and_prompt(image, image_url, prompt, temperature=0.3, max_toke
             return f"Error loading image from URL: {e}"
     else:
         return "Please provide either an image or an image URL."
-    
-    # Format message with the aya-vision chat template
+
+    # Format the message using the Aya Vision chat template
     messages = [
-        {"role": "user",
-         "content": [
-           {"type": "image", "source": img},
-            {"type": "text", "text": prompt},
-        ]},
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "source": img},
+                {"type": "text", "text": prompt},
+            ],
+        },
     ]
 
-    # Process input
     try:
         inputs = processor.apply_chat_template(
-            messages, 
-            padding=True, 
-            add_generation_prompt=True, 
-            tokenize=True, 
-            return_dict=True, 
+            messages,
+            padding=True,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
             return_tensors="pt"
         ).to(model.device)
-        
-        # Generate response
+
         gen_tokens = model.generate(
-            **inputs, 
-            max_new_tokens=int(max_tokens), 
-            do_sample=True, 
+            **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)
+        response = processor.tokenizer.decode(
+            gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True
+        )
         return response
     except Exception as e:
         return f"Error generating response: {e}"
 
-# Define example inputs
+# 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 Gradio application
+# Create the Gradio application
 with gr.Blocks(title="Aya Vision 8B Demo") as demo:
     gr.Markdown("# Aya Vision 8B Model Demo")
     gr.Markdown("""
@@ -97,36 +97,32 @@ with gr.Blocks(title="Aya Vision 8B Demo") as demo:
     - Visual reasoning
     - Question answering
     - Support for 23 languages
-    
+
     Upload an image or provide a URL, and enter a prompt to get started!
     """)
     
-    with gr.Row():
-        with gr.Column():
-            load_button = gr.Button("Load Model", variant="primary")
-            status = gr.Textbox(label="Model Status", placeholder="Model not loaded yet. Click 'Load Model' to start.")
-            
-            gr.Markdown("### Upload an image or provide an image 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)
-            
-            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)
-                
-            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)
-    
-    # Add examples section
+    # Display model loading status
+    gr.Markdown(f"**Model Status:** {model_status}")
+
+    gr.Markdown("### Upload an image or provide an image 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)
+    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)
+
+    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,
@@ -135,20 +131,10 @@ with gr.Blocks(title="Aya Vision 8B Demo") as demo:
         fn=process_image_and_prompt
     )
 
-    # Set up tab switching logic - hide appropriate inputs depending on tab
-    def update_image_tab():
-        return {image_url_input: gr.update(visible=False), image_input: gr.update(visible=True)}
-    
-    def update_url_tab():
-        return {image_url_visible: gr.update(visible=True), image_input_url: gr.update(visible=False)}
-    
-    # Define button click behavior
-    load_button.click(load_model, inputs=None, outputs=status)
-    
     # 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)
-    
+
     generate_button.click(
         generate_response,
         inputs=[image_input, image_url_visible, prompt, temperature, max_tokens],