Allow nested transformer block names for quantization, lora targeting, quantizing

extensions_built_in/diffusion_models/hidream/hidream_o1_model.py

@@ -127,6 +127,7 @@ class HidreamO1Model(BaseModel):
         super().__init__(
             device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
         )
+        self.use_old_lokr_format = False
         self.is_flow_matching = True
         self.is_transformer = True
         self.target_lora_modules = ["Qwen3VLForConditionalGeneration"]
@@ -524,7 +525,7 @@ class HidreamO1Model(BaseModel):
         return self.arch
 
     def get_transformer_block_names(self) -> Optional[List[str]]:
-        return ["layers"]
+        return ["model.language_model.layers"]
 
     def convert_lora_weights_before_save(self, state_dict):
         new_sd = {}

jobs/process/BaseSDTrainProcess.py

@@ -2135,7 +2135,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
                         compiled_block_count = 0
 
                         for attr_name in BLOCK_LIST_ATTRS:
-                            block_list = getattr(inner_unet, attr_name, None)
+                            # attr_name may be a dotted path for models that nest their
+                            # blocks (e.g. hidream_o1's "model.language_model.layers").
+                            block_list = inner_unet
+                            for part in attr_name.split('.'):
+                                block_list = getattr(block_list, part, None)
+                                if block_list is None:
+                                    break
 
                             if block_list is None:
                                 continue

toolkit/lora_special.py

@@ -370,7 +370,10 @@ class LoRASpecialNetwork(ToolkitNetworkMixin, LoRANetwork):
                                 transformer_block_names = base_model.get_transformer_block_names()
                             
                             if transformer_block_names is not None:
-                                if not any([name in lora_name for name in transformer_block_names]):
+                                # match against clean_name (dotted) so block names can be
+                                # dotted paths (e.g. "model.language_model.layers"); lora_name
+                                # has dots replaced with "$$"/"_" and wouldn't match.
+                                if not any([block_name in clean_name for block_name in transformer_block_names]):
                                     skip = True
                             else:
                                 if self.is_pixart:

toolkit/util/quantize.py

@@ -297,7 +297,13 @@ def quantize_model(
         all_blocks: List[torch.nn.Module] = []
         transformer_block_names = base_model.get_transformer_block_names()
         for name in transformer_block_names:
-            block_list = getattr(model_to_quantize, name, None)
+            # name may be a dotted path for models that nest their blocks
+            # (e.g. hidream_o1's "model.language_model.layers").
+            block_list = model_to_quantize
+            for part in name.split('.'):
+                block_list = getattr(block_list, part, None)
+                if block_list is None:
+                    break
             if block_list is not None:
                 all_blocks += list(block_list)
         base_model.print_and_status_update(