techfreakworm commited on
Commit
be7dc01
·
unverified ·
1 Parent(s): 1494a15

fix(spaces): module-level @spaces.GPU so ZeroGPU's startup detector finds it

Browse files

ZeroGPU rejected the previous build with "No @spaces.GPU function detected
during startup". The decorator was being applied dynamically inside
submit() per-call, which the startup analyzer can't see.

Refactored to a module-level _execute_workflow(executor, workflow,
output_ids) decorated at import time with spaces.GPU(duration=300)
(Pro per-call cap). _worker now calls this synchronous helper from its
thread; the per-call gpu_duration plumbing is no longer used because
ZeroGPU dispatch is bound at decoration time, not call time.

Files changed (1) hide show
  1. backend.py +35 -19
backend.py CHANGED
@@ -53,6 +53,34 @@ def _on_spaces() -> bool:
53
  return bool(os.environ.get("SPACES_ZERO_GPU"))
54
 
55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
  class _StubServer:
57
  """Minimal stub matching the surface ComfyUI's PromptExecutor expects."""
58
 
@@ -239,7 +267,7 @@ class ComfyUILibraryBackend:
239
  return f"ComfyUILibraryBackend(comfy_dir={self._comfy_dir!r})"
240
 
241
  async def submit(
242
- self, mode: str, workflow: dict, gpu_duration: int = 120
243
  ) -> AsyncIterator[Any]:
244
  """Run a workflow end-to-end. Yields Download/Progress/Output/Error events."""
245
  # Pre-flight: ensure all model files exist.
@@ -309,17 +337,11 @@ class ComfyUILibraryBackend:
309
  # Use the public setter; it writes the same global the
310
  # ProgressBar class reads, but is the documented API.
311
  comfy.utils.set_progress_bar_global_hook(_hook)
312
- self._executor.execute(
313
- workflow,
314
- prompt_id="ltx23-aio",
315
- extra_data={"client_id": "ltx23-aio"},
316
- execute_outputs=output_ids,
317
- )
318
- # PromptExecutor stores per-node UI info in history_result["outputs"]
319
- # after execute_async. Each entry mirrors what the JS frontend
320
- # would receive — including SaveVideo's "filenames"/"video" lists
321
- # that point at the saved file inside ComfyUI's output dir.
322
- hist = getattr(self._executor, "history_result", {}) or {}
323
  outs = hist.get("outputs") or {}
324
  video_path = _first_video_path(list(outs.values())) or ""
325
  _push(OutputEvent(video_path=video_path))
@@ -338,13 +360,7 @@ class ComfyUILibraryBackend:
338
  _free_memory()
339
  _push(None) # sentinel: stop the consumer
340
 
341
- if _on_spaces():
342
- import spaces
343
-
344
- execute = spaces.GPU(duration=gpu_duration)(_worker)
345
- thread = threading.Thread(target=execute, daemon=True)
346
- else:
347
- thread = threading.Thread(target=_worker, daemon=True)
348
  thread.start()
349
 
350
  while True:
 
53
  return bool(os.environ.get("SPACES_ZERO_GPU"))
54
 
55
 
56
+ try:
57
+ import spaces # type: ignore
58
+ except ImportError:
59
+ spaces = None # type: ignore[assignment]
60
+
61
+
62
+ def _identity(fn):
63
+ return fn
64
+
65
+
66
+ # ZeroGPU's startup detector scans loaded modules for spaces.GPU-wrapped
67
+ # functions. The decorator must be applied at module load time — runtime
68
+ # wrapping inside a request handler isn't detected. Pro tier per-call cap is
69
+ # 300s, so we use that ceiling and let modes finish whenever they finish.
70
+ _GPU = spaces.GPU(duration=300) if (spaces is not None and _on_spaces()) else _identity
71
+
72
+
73
+ @_GPU
74
+ def _execute_workflow(executor: Any, workflow: dict, output_ids: list[str]) -> dict:
75
+ executor.execute(
76
+ workflow,
77
+ prompt_id="ltx23-aio",
78
+ extra_data={"client_id": "ltx23-aio"},
79
+ execute_outputs=output_ids,
80
+ )
81
+ return getattr(executor, "history_result", {}) or {}
82
+
83
+
84
  class _StubServer:
85
  """Minimal stub matching the surface ComfyUI's PromptExecutor expects."""
86
 
 
267
  return f"ComfyUILibraryBackend(comfy_dir={self._comfy_dir!r})"
268
 
269
  async def submit(
270
+ self, mode: str, workflow: dict, gpu_duration: int = 300
271
  ) -> AsyncIterator[Any]:
272
  """Run a workflow end-to-end. Yields Download/Progress/Output/Error events."""
273
  # Pre-flight: ensure all model files exist.
 
337
  # Use the public setter; it writes the same global the
338
  # ProgressBar class reads, but is the documented API.
339
  comfy.utils.set_progress_bar_global_hook(_hook)
340
+ # _execute_workflow is module-level and decorated with
341
+ # @spaces.GPU(duration=300) on Spaces — that's what makes the
342
+ # heavy compute run on a borrowed H200. Off-Spaces it's a
343
+ # plain call.
344
+ hist = _execute_workflow(self._executor, workflow, output_ids)
 
 
 
 
 
 
345
  outs = hist.get("outputs") or {}
346
  video_path = _first_video_path(list(outs.values())) or ""
347
  _push(OutputEvent(video_path=video_path))
 
360
  _free_memory()
361
  _push(None) # sentinel: stop the consumer
362
 
363
+ thread = threading.Thread(target=_worker, daemon=True)
 
 
 
 
 
 
364
  thread.start()
365
 
366
  while True: