desiree commited on
Commit
2cc50ab
·
verified ·
1 Parent(s): 4c3ee66

Keep Wan 2.1-1.3B selectable in small lane (prefer='wan21_t2v_1_3b')

Browse files
Files changed (1) hide show
  1. 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). The router picks the cheapest capable one: CogVideoX-2b
21
- # (fp16, ~5.3 GB VRAM, 6 steps) beats Wan 2.1 T2V-1.3B (bf16, ~8.2 GB) for drafts.
 
22
  "video": CogVideoRenderer("THUDM/CogVideoX-2b"),
23
  "image": MockImageRenderer(),
24
  "audio": MockAudioRenderer(),
25
  },
26
  "production": {
27
- # placeholder quality lane: loads on first use, needs ~80 GB VRAM (Phase 10 hardware)
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