goumsss Claude Sonnet 4.6 commited on
Commit
885e6f8
·
1 Parent(s): 0a993a9

Fix: restore @spaces.GPU decorator required by ZeroGPU

Browse files

Decorator was accidentally dropped during model refactor.
ZeroGPU fails at startup with "No @spaces.GPU function detected"
without it.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Files changed (1) hide show
  1. image_generator.py +9 -3
image_generator.py CHANGED
@@ -167,6 +167,12 @@ def _generate(animals: list[str], places: list[str]):
167
  # Public API — two versions depending on environment
168
  # ---------------------------------------------------------------------------
169
 
170
- def generate_reward_image(animals: list[str], places: list[str]):
171
- """Generate reward image — works locally (MPS/CPU) and on HF Spaces (persistent GPU)."""
172
- return _generate(animals, places)
 
 
 
 
 
 
 
167
  # Public API — two versions depending on environment
168
  # ---------------------------------------------------------------------------
169
 
170
+ if IS_HF_SPACE:
171
+ @spaces.GPU(duration=60)
172
+ def generate_reward_image(animals: list[str], places: list[str]):
173
+ """Generate reward image on HF Spaces ZeroGPU."""
174
+ return _generate(animals, places)
175
+ else:
176
+ def generate_reward_image(animals: list[str], places: list[str]):
177
+ """Generate reward image locally (MPS/CPU)."""
178
+ return _generate(animals, places)