Fix: restore @spaces.GPU decorator required by ZeroGPU
Browse filesDecorator 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>
- 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 |
-
|
| 171 |
-
|
| 172 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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)
|