doubility123 commited on
Commit
5c9ea2e
·
1 Parent(s): 90b258a

add huggingface space mode

Browse files
Files changed (1) hide show
  1. deepseek_vl2/serve/inference.py +5 -1
deepseek_vl2/serve/inference.py CHANGED
@@ -19,6 +19,7 @@
19
 
20
  from threading import Thread
21
  from typing import List
 
22
 
23
  import torch
24
  import transformers
@@ -40,7 +41,9 @@ def load_model(model_path, dtype=torch.bfloat16):
40
  vl_gpt: DeepseekVLV2ForCausalLM = AutoModelForCausalLM.from_pretrained(
41
  model_path, trust_remote_code=True, torch_dtype=dtype
42
  )
43
- vl_gpt = vl_gpt.cuda().eval()
 
 
44
  return tokenizer, vl_gpt, vl_chat_processor
45
 
46
 
@@ -87,6 +90,7 @@ class StoppingCriteriaSub(StoppingCriteria):
87
 
88
 
89
  @torch.inference_mode()
 
90
  def deepseek_generate(
91
  conversations: list,
92
  vl_gpt: torch.nn.Module,
 
19
 
20
  from threading import Thread
21
  from typing import List
22
+ import spaces
23
 
24
  import torch
25
  import transformers
 
41
  vl_gpt: DeepseekVLV2ForCausalLM = AutoModelForCausalLM.from_pretrained(
42
  model_path, trust_remote_code=True, torch_dtype=dtype
43
  )
44
+ if torch.cuda.is_available():
45
+ vl_gpt = vl_gpt.cuda()
46
+
47
  return tokenizer, vl_gpt, vl_chat_processor
48
 
49
 
 
90
 
91
 
92
  @torch.inference_mode()
93
+ @spaces.GPU(duration=120)
94
  def deepseek_generate(
95
  conversations: list,
96
  vl_gpt: torch.nn.Module,