arad1367 commited on
Commit
0c9f542
1 Parent(s): 34d3a3a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -3,6 +3,9 @@ import gradio as gr
3
  import torch
4
  from PIL import Image
5
  from transformers import AutoModel, AutoTokenizer
 
 
 
6
 
7
  # Load the model and tokenizer
8
  model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
@@ -11,6 +14,7 @@ tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_
11
  model.eval()
12
 
13
  # Define a function to generate a response
 
14
  def generate_response(image, question):
15
  msgs = [{'role': 'user', 'content': question}]
16
  res = model.chat(
 
3
  import torch
4
  from PIL import Image
5
  from transformers import AutoModel, AutoTokenizer
6
+ import spaces
7
+
8
+ device="cuda"
9
 
10
  # Load the model and tokenizer
11
  model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
 
14
  model.eval()
15
 
16
  # Define a function to generate a response
17
+ @spaces.GPU
18
  def generate_response(image, question):
19
  msgs = [{'role': 'user', 'content': question}]
20
  res = model.chat(