Phase 2

toolkit/util/get_model.py

@@ -41,10 +41,34 @@ def get_all_models() -> List[BaseModel]:
     return all_model_classes
 
 
+# archs the legacy StableDiffusion monolith still serves (see the arch
+# normalization in toolkit/config_modules.py)
+LEGACY_ARCHS = {
+    "sd1",
+    "sd2",
+    "sd3",
+    "sdxl",
+    "pixart",
+    "pixart_sigma",
+    "auraflow",
+    "flux",
+    "lumina2",
+    "vega",
+    "ssd",
+}
+
+
 def get_model_class(config: ModelConfig):
     all_models = get_all_models()
     for ModelClass in all_models:
         if ModelClass.arch == config.arch:
             return ModelClass
-    # default to the legacy model
-    return StableDiffusion
+    if config.arch in LEGACY_ARCHS:
+        return StableDiffusion
+    # a typo'd or unregistered arch used to silently fall back to SD1; error
+    # instead (a broken extension import also lands here — its error was
+    # printed during get_all_models)
+    known = sorted({m.arch for m in all_models if m.arch} | LEGACY_ARCHS)
+    raise ValueError(
+        f"Unknown model arch {config.arch!r}. Known archs: {', '.join(known)}"
+    )

toolkit/util/quantize.py

@@ -429,9 +429,10 @@ def quantize_model(
         # quantize model the original way without an accuracy recovery adapter
         # move and quantize only certain pieces at a time.
         quantization_type = get_qtype(base_model.model_config.qtype)
+        quantize_kwargs = base_model.model_config.quantize_kwargs or {}
         # all_blocks = list(model_to_quantize.transformer_blocks)
         all_blocks: List[torch.nn.Module] = []
-        transformer_block_names = base_model.get_transformer_block_names()
+        transformer_block_names = base_model.get_transformer_block_names() or []
         for name in transformer_block_names:
             # name may be a dotted path for models that nest their blocks
             # (e.g. hidream_o1's "model.language_model.layers").
@@ -456,7 +457,12 @@ def quantize_model(
             block.to(base_model.device_torch, dtype=base_model.torch_dtype, non_blocking=True)
             # exclude patterns with a leading wildcard (e.g. "*adaln_proj*")
             # also apply inside blocks, where names are block-relative
-            quantize(block, weights=quantization_type, exclude=exclude_modules)
+            quantize(
+                block,
+                weights=quantization_type,
+                exclude=exclude_modules,
+                **quantize_kwargs,
+            )
             freeze(block)
             # NOT non_blocking: an async D2H allocates the cpu destination in pinned
             # memory, which the caching host allocator keeps forever (with power-of-2
@@ -472,5 +478,10 @@ def quantize_model(
         # device without having to move the transformer blocks to the device first
         base_model.print_and_status_update(" - quantizing extras")
         # model_to_quantize.to(base_model.device_torch, dtype=base_model.torch_dtype)
-        quantize(model_to_quantize, weights=quantization_type, exclude=exclude_modules)
+        quantize(
+            model_to_quantize,
+            weights=quantization_type,
+            exclude=exclude_modules,
+            **quantize_kwargs,
+        )
         freeze(model_to_quantize)