Georg Claude Sonnet 4.5 commited on
Commit
4fd63bf
·
1 Parent(s): e93c3d3

Add GPU warmup function for ZeroGPU detection

Browse files

- Create warmup() function decorated with @spaces.GPU(duration=10)
- Add @spaces.GPU(duration=120) back to initialize_model()
- Call warmup() on gradio_app.load() to ensure GPU detected at startup
- Remove @spaces.GPU from wrapper functions (they call GPU methods internally)

ZeroGPU requires at least one GPU function invoked during startup.

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

Files changed (1) hide show
  1. app.py +13 -2
app.py CHANGED
@@ -38,6 +38,7 @@ class FoundationPoseInference:
38
  self.tracked_objects = {}
39
  self.use_real_model = USE_REAL_MODEL
40
 
 
41
  def initialize_model(self):
42
  """Initialize the FoundationPose model on GPU."""
43
  if self.initialized:
@@ -262,8 +263,16 @@ async def api_estimate(request: EstimateRequest):
262
  raise HTTPException(status_code=500, detail=str(e))
263
 
264
 
 
 
 
 
 
 
 
 
 
265
  # Gradio wrapper functions
266
- @spaces.GPU(duration=120)
267
  def gradio_initialize(object_id: str, reference_files: List, fx: float, fy: float, cx: float, cy: float):
268
  """Gradio wrapper for object initialization."""
269
  try:
@@ -307,7 +316,6 @@ def gradio_initialize(object_id: str, reference_files: List, fx: float, fy: floa
307
  return f"Error: {str(e)}"
308
 
309
 
310
- @spaces.GPU(duration=30)
311
  def gradio_estimate(object_id: str, query_image: np.ndarray, fx: float, fy: float, cx: float, cy: float):
312
  """Gradio wrapper for pose estimation."""
313
  try:
@@ -484,6 +492,9 @@ with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as grad
484
  See the [API documentation](https://huggingface.co/spaces/gpue/foundationpose) for details.
485
  """)
486
 
 
 
 
487
 
488
  # Mount Gradio to FastAPI
489
  app = gr.mount_gradio_app(app, gradio_app, path="/")
 
38
  self.tracked_objects = {}
39
  self.use_real_model = USE_REAL_MODEL
40
 
41
+ @spaces.GPU(duration=120)
42
  def initialize_model(self):
43
  """Initialize the FoundationPose model on GPU."""
44
  if self.initialized:
 
263
  raise HTTPException(status_code=500, detail=str(e))
264
 
265
 
266
+ # Warmup function to ensure ZeroGPU detects GPU usage
267
+ @spaces.GPU(duration=10)
268
+ def warmup():
269
+ """Warmup function to initialize GPU context for ZeroGPU."""
270
+ logger.info("Warming up GPU for ZeroGPU...")
271
+ pose_estimator.initialize_model()
272
+ return "✓ GPU initialized"
273
+
274
+
275
  # Gradio wrapper functions
 
276
  def gradio_initialize(object_id: str, reference_files: List, fx: float, fy: float, cx: float, cy: float):
277
  """Gradio wrapper for object initialization."""
278
  try:
 
316
  return f"Error: {str(e)}"
317
 
318
 
 
319
  def gradio_estimate(object_id: str, query_image: np.ndarray, fx: float, fy: float, cx: float, cy: float):
320
  """Gradio wrapper for pose estimation."""
321
  try:
 
492
  See the [API documentation](https://huggingface.co/spaces/gpue/foundationpose) for details.
493
  """)
494
 
495
+ # Warmup on load to ensure ZeroGPU detects GPU usage
496
+ gradio_app.load(warmup, outputs=None)
497
+
498
 
499
  # Mount Gradio to FastAPI
500
  app = gr.mount_gradio_app(app, gradio_app, path="/")