resakemal commited on
Commit
d890a16
·
1 Parent(s): 7f30dac

ZeroGPU fix attempt

Browse files
Files changed (1) hide show
  1. model/backend.py +9 -7
model/backend.py CHANGED
@@ -15,12 +15,13 @@ import os, json, re, urllib.request
15
 
16
  # ZeroGPU: decorate GPU-using functions with @_GPU so HF Spaces can allocate
17
  # the GPU on demand. Falls back to a no-op when running locally (no `spaces` pkg).
18
- try:
19
- import spaces
20
- _GPU = spaces.GPU
21
- except ImportError:
22
- def _GPU(fn=None, *, duration=60):
23
- return fn if fn is not None else (lambda f: f)
 
24
 
25
  # is_torch_fx_available fix
26
  import transformers.utils.import_utils
@@ -188,7 +189,8 @@ def _get_hf_pipe():
188
  # Stateless (one-turn only), strictly prevents thinking tokens per Qwen3 docs.
189
  _NO_THINK_SUFFIX = [{"role": "assistant", "content": "<think>\n\n</think>\n\n"}]
190
 
191
- @_GPU(duration=120)
 
192
  def _hf_generate(messages, max_new_tokens=512, temperature=0.6) -> str:
193
  pipe = _get_hf_pipe()
194
  full_messages = messages + _NO_THINK_SUFFIX
 
15
 
16
  # ZeroGPU: decorate GPU-using functions with @_GPU so HF Spaces can allocate
17
  # the GPU on demand. Falls back to a no-op when running locally (no `spaces` pkg).
18
+ import spaces
19
+ # try:
20
+ # import spaces
21
+ # # _GPU = spaces.GPU
22
+ # except ImportError:
23
+ # def _GPU(fn=None, *, duration=60):
24
+ # return fn if fn is not None else (lambda f: f)
25
 
26
  # is_torch_fx_available fix
27
  import transformers.utils.import_utils
 
189
  # Stateless (one-turn only), strictly prevents thinking tokens per Qwen3 docs.
190
  _NO_THINK_SUFFIX = [{"role": "assistant", "content": "<think>\n\n</think>\n\n"}]
191
 
192
+ # @_GPU(duration=120)
193
+ @spaces.GPU(duration=120)
194
  def _hf_generate(messages, max_new_tokens=512, temperature=0.6) -> str:
195
  pipe = _get_hf_pipe()
196
  full_messages = messages + _NO_THINK_SUFFIX