File size: 1,405 Bytes
1e3ce4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """Op dispatch registry.
The cursor core defines the op contract; backends (in studio/backends/) register
concrete implementations via register_backend. Keeps pixel_cursor/ free of model
imports — torch/diffusers/MotionClone never enter the core package.
"""
from __future__ import annotations
from typing import Callable, Any
_REGISTRY: dict[tuple[str, str], Callable] = {}
def register_backend(op_name: str, backend_name: str, fn: Callable) -> None:
_REGISTRY[(op_name, backend_name)] = fn
def list_backends(op_name: str | None = None) -> list[tuple[str, str]]:
if op_name is None:
return sorted(_REGISTRY.keys())
return sorted(k for k in _REGISTRY if k[0] == op_name)
def _dispatch(op_name: str, backend: str, cursor, *args, **kwargs) -> Any:
if backend == "auto":
for (op, _name), fn in _REGISTRY.items():
if op == op_name:
return fn(cursor, *args, **kwargs)
raise NotImplementedError(
f"No backend registered for op {op_name!r}. "
f"Register one via pixel_cursor.ops.register_backend()."
)
fn = _REGISTRY.get((op_name, backend))
if fn is None:
available = sorted({n for (o, n) in _REGISTRY if o == op_name})
raise NotImplementedError(
f"No backend {backend!r} for op {op_name!r}; available: {available}"
)
return fn(cursor, *args, **kwargs)
|