rydlrKE commited on
Commit
c6cdc8c
·
verified ·
1 Parent(s): f387ec5

ZeroGPU-safe startup device + llm2vec tensor warning cleanup

Browse files
Files changed (1) hide show
  1. kimodo/demo/app.py +8 -1
kimodo/demo/app.py CHANGED
@@ -16,6 +16,7 @@ import viser
16
  from kimodo.assets import DEMO_ASSETS_ROOT
17
  from kimodo.model.load_model import load_model
18
  from kimodo.model.registry import resolve_model_name
 
19
  from kimodo.skeleton import SkeletonBase, SOMASkeleton30
20
  from kimodo.tools import load_json
21
  from kimodo.viz import viser_utils
@@ -60,7 +61,13 @@ logging.getLogger("websockets.asyncio.server").setLevel(logging.CRITICAL)
60
 
61
  class Demo:
62
  def __init__(self, default_model_name: str = DEFAULT_MODEL):
63
- self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
 
 
 
 
64
  print(f"Using device: {self.device}")
65
  self.models: dict[str, ModelBundle] = {}
66
  resolved = resolve_model_name(default_model_name, "Kimodo")
 
16
  from kimodo.assets import DEMO_ASSETS_ROOT
17
  from kimodo.model.load_model import load_model
18
  from kimodo.model.registry import resolve_model_name
19
+ from kimodo.runtime.device import select_runtime_device
20
  from kimodo.skeleton import SkeletonBase, SOMASkeleton30
21
  from kimodo.tools import load_json
22
  from kimodo.viz import viser_utils
 
61
 
62
  class Demo:
63
  def __init__(self, default_model_name: str = DEFAULT_MODEL):
64
+ # In hosted HF runtimes (including ZeroGPU), touching CUDA too early can
65
+ # crash startup before queue-managed inference starts.
66
+ requested_device = os.getenv("KIMODO_DEVICE")
67
+ running_in_space = bool(os.getenv("SPACE_ID")) or os.getenv("SYSTEM", "").strip().lower() == "spaces"
68
+ if requested_device is None and (HF_MODE or running_in_space):
69
+ requested_device = "cpu"
70
+ self.device = select_runtime_device(requested=requested_device)
71
  print(f"Using device: {self.device}")
72
  self.models: dict[str, ModelBundle] = {}
73
  resolved = resolve_model_name(default_model_name, "Kimodo")