Update app.py

app.py

@@ -4,21 +4,25 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 import requests
 import gradio as gr
+import spaces  # Import Hugging Face Spaces package
 
 # Load model and tokenizer
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = AutoModelForCausalLM.from_pretrained(
-    'scb10x/llama-3-typhoon-v1.5-8b-instruct-vision-preview',
-    revision='main',  # or specify a commit hash if needed
-    torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
-    device_map='auto',
-    trust_remote_code=True
-).to(device)
+model_name = 'scb10x/llama-3-typhoon-v1.5-8b-instruct-vision-preview'
 
-tokenizer = AutoTokenizer.from_pretrained(
-    'scb10x/llama-3-typhoon-v1.5-8b-instruct-vision-preview',
-    trust_remote_code=True
-)
+@spaces.GPU(duration=60)  # Decorate the function to dynamically request and release GPU
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
+        device_map='auto',
+        trust_remote_code=True
+    )
+    return model
+
+model = load_model()
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 def prepare_inputs(text, image, device='cuda'):
     messages = [
@@ -38,7 +42,7 @@ def prepare_inputs(text, image, device='cuda'):
 
     return input_ids, attention_mask
 
-# Inference function
+@spaces.GPU(duration=60)  # Decorate the function for GPU use
 def predict(prompt, img_url):
     try:
         image = Image.open(requests.get(img_url, stream=True).raw)