Update app.py

app.py

@@ -5,6 +5,7 @@ import gradio as gr
 from threading import Thread
 from PIL import Image
 import subprocess
+import spaces  # Add this import
 
 # Install flash-attention
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -46,6 +47,7 @@ vision_model = AutoModelForCausalLM.from_pretrained(
 vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
 
 # Helper functions
+@spaces.GPU  # Add this decorator
 def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_tokens=1024, top_p=1.0, top_k=20):
     conversation = [{"role": "system", "content": system_prompt}]
     for prompt, answer in history:
@@ -78,6 +80,7 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
         buffer += new_text
         yield buffer
 
+@spaces.GPU  # Add this decorator
 def process_vision_query(image, text_input):
     prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n"
     image = Image.fromarray(image).convert("RGB")