Kaiky Maia commited on
Commit
cd53ece
·
verified ·
1 Parent(s): 680bfe9

ZeroGPU: load student on CPU at startup with cuda.is_available guarded False

Browse files
Files changed (1) hide show
  1. server_app.py +12 -8
server_app.py CHANGED
@@ -36,11 +36,13 @@ def viz() -> dict:
36
  @app.api(name="probe", stream_every=0.1)
37
  @spaces.GPU
38
  def probe(text: str) -> dict: # annotation = type of each streamed chunk
39
- if not text.strip() or not RT.model_ready:
40
- yield {"error": "model not ready" if text.strip() else "empty prompt",
41
- "done": True}
 
 
 
42
  return
43
- _boot.to_device(RT)
44
  partial, gates = "", []
45
  for partial, gates in _boot.stream(RT, text):
46
  yield {"text": partial, "gates": gates, "done": False}
@@ -54,11 +56,13 @@ def probe(text: str) -> dict: # annotation = type of each streamed chunk
54
  @app.api(name="arena", stream_every=0.1)
55
  @spaces.GPU
56
  def arena(text: str) -> dict: # annotation = type of each streamed chunk
57
- if not text.strip() or not RT.model_ready:
58
- yield {"error": "model not ready" if text.strip() else "empty prompt",
59
- "done": True}
 
 
 
60
  return
61
- _boot.to_device(RT)
62
  base = deku = ""
63
  gates: list[float] = []
64
  for base, deku, gates in _boot.stream_pair(RT, text):
 
36
  @app.api(name="probe", stream_every=0.1)
37
  @spaces.GPU
38
  def probe(text: str) -> dict: # annotation = type of each streamed chunk
39
+ if not text.strip():
40
+ yield {"error": "empty prompt", "done": True}
41
+ return
42
+ _boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU)
43
+ if not RT.model_ready:
44
+ yield {"error": RT.load_error or "model not ready", "done": True}
45
  return
 
46
  partial, gates = "", []
47
  for partial, gates in _boot.stream(RT, text):
48
  yield {"text": partial, "gates": gates, "done": False}
 
56
  @app.api(name="arena", stream_every=0.1)
57
  @spaces.GPU
58
  def arena(text: str) -> dict: # annotation = type of each streamed chunk
59
+ if not text.strip():
60
+ yield {"error": "empty prompt", "done": True}
61
+ return
62
+ _boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU)
63
+ if not RT.model_ready:
64
+ yield {"error": RT.load_error or "model not ready", "done": True}
65
  return
 
66
  base = deku = ""
67
  gates: list[float] = []
68
  for base, deku, gates in _boot.stream_pair(RT, text):