Keep Wan 2.1-1.3B selectable in small lane (prefer='wan21_t2v_1_3b')
Browse files- spi/renderers/registry.py +12 -3
spi/renderers/registry.py
CHANGED
|
@@ -17,20 +17,26 @@ DEFAULT_REGISTRIES: dict[str, dict[str, Renderer]] = {
|
|
| 17 |
"audio": MockAudioRenderer(),
|
| 18 |
},
|
| 19 |
"small": {
|
| 20 |
-
# real backends (Phase 2).
|
| 21 |
-
#
|
|
|
|
| 22 |
"video": CogVideoRenderer("THUDM/CogVideoX-2b"),
|
| 23 |
"image": MockImageRenderer(),
|
| 24 |
"audio": MockAudioRenderer(),
|
| 25 |
},
|
| 26 |
"production": {
|
| 27 |
-
#
|
| 28 |
"video": WanRenderer("Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
|
| 29 |
"image": MockImageRenderer(),
|
| 30 |
"audio": MockAudioRenderer(),
|
| 31 |
},
|
| 32 |
}
|
| 33 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 34 |
|
| 35 |
class ModelRouter:
|
| 36 |
def __init__(self, mode: str = "mock") -> None:
|
|
@@ -41,6 +47,9 @@ class ModelRouter:
|
|
| 41 |
kind: {renderer.name: renderer}
|
| 42 |
for kind, renderer in DEFAULT_REGISTRIES[mode].items()
|
| 43 |
}
|
|
|
|
|
|
|
|
|
|
| 44 |
|
| 45 |
def register(self, kind: str, renderer: Renderer) -> None:
|
| 46 |
self._lanes.setdefault(kind, {})[renderer.name] = renderer
|
|
|
|
| 17 |
"audio": MockAudioRenderer(),
|
| 18 |
},
|
| 19 |
"small": {
|
| 20 |
+
# real backends (Phase 2). Cheap-first route picks CogVideoX-2b (fp16, ~5.3 GB,
|
| 21 |
+
# 6 steps); Wan 2.1 T2V-1.3B (bf16, ~8.2 GB, ~29 GB download incl. fp32 UMT5-XXL)
|
| 22 |
+
# stays selectable via prefer="wan21_t2v_1_3b".
|
| 23 |
"video": CogVideoRenderer("THUDM/CogVideoX-2b"),
|
| 24 |
"image": MockImageRenderer(),
|
| 25 |
"audio": MockAudioRenderer(),
|
| 26 |
},
|
| 27 |
"production": {
|
| 28 |
+
# Phase 10 quality lane placeholder: loads on first use, needs ~80 GB VRAM
|
| 29 |
"video": WanRenderer("Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
|
| 30 |
"image": MockImageRenderer(),
|
| 31 |
"audio": MockAudioRenderer(),
|
| 32 |
},
|
| 33 |
}
|
| 34 |
|
| 35 |
+
# adapters built + tested in Phase 2 that are not the default route in any lane
|
| 36 |
+
ADDITIONAL_SMALL_RENDERERS = {
|
| 37 |
+
"wan21_t2v_1_3b": WanRenderer("Wan-AI/Wan2.1-T2V-1.3B-Diffusers"),
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
|
| 41 |
class ModelRouter:
|
| 42 |
def __init__(self, mode: str = "mock") -> None:
|
|
|
|
| 47 |
kind: {renderer.name: renderer}
|
| 48 |
for kind, renderer in DEFAULT_REGISTRIES[mode].items()
|
| 49 |
}
|
| 50 |
+
if mode == "small": # tested Phase 2 adapters remain reachable via prefer=
|
| 51 |
+
for name, renderer in ADDITIONAL_SMALL_RENDERERS.items():
|
| 52 |
+
self.register("video", renderer)
|
| 53 |
|
| 54 |
def register(self, kind: str, renderer: Renderer) -> None:
|
| 55 |
self._lanes.setdefault(kind, {})[renderer.name] = renderer
|