shb777 commited on
Commit
9dc5c64
·
verified ·
1 Parent(s): 3f01084

Update app.py

Browse files

Update space for ZeroGPU as per [0dac0d94a8898ca9c104274ea91258fc534dc362](https://huggingface.co/spaces/victor/Granite-Vision-3.1-2B/commit/0dac0d94a8898ca9c104274ea91258fc534dc362)

Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import random
2
  import torch
3
  import gradio as gr
@@ -16,6 +17,7 @@ def get_text_from_content(content):
16
  texts.append("[Image]")
17
  return " ".join(texts)
18
 
 
19
  def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversation):
20
  if conversation is None:
21
  conversation = []
@@ -39,7 +41,7 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
39
  tokenize=True,
40
  return_dict=True,
41
  return_tensors="pt"
42
- ).to("cpu")
43
 
44
  torch.manual_seed(random.randint(0, 10000))
45
 
 
1
+ import spaces
2
  import random
3
  import torch
4
  import gradio as gr
 
17
  texts.append("[Image]")
18
  return " ".join(texts)
19
 
20
+ @spaces.GPU
21
  def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversation):
22
  if conversation is None:
23
  conversation = []
 
41
  tokenize=True,
42
  return_dict=True,
43
  return_tensors="pt"
44
+ ).to("cuda")
45
 
46
  torch.manual_seed(random.randint(0, 10000))
47