archive_model_dtypes called unconditionally in BaseModel.__init__ breaks custom nodes using disable_unet_model_creation

[ivanprado]

Custom Node Testing

• I have tried disabling custom nodes and the issue persists (see how to disable custom nodes if you need help)

Expected Behavior

When a custom node extends comfy.model_base.BaseModel and sets unet_config["disable_unet_model_creation"] = True, the BaseModel.init method should skip all operations related to self.diffusion_model since the child class will create it manually after calling super().init().

Actual Behavior

The comfy.model_management.archive_model_dtypes(self.diffusion_model) call in BaseModel.init is placed outside the if not unet_config.get("disable_unet_model_creation", False): block, causing an AttributeError because self.diffusion_model does not exist when disable_unet_model_creation=True.

Steps to Reproduce

Create a custom node that extends comfy.model_base.BaseModel
Set unet_config["disable_unet_model_creation"] = True in the model config
Call super().init() and then set self.diffusion_model manually afterward
Run a workflow using this custom node
Error occurs during model initialization

Minimal reproduction code:

class CustomModelConfig:
    def __init__(self, dtype):
        self.unet_config = {}
        self.unet_config["disable_unet_model_creation"] = True
        # ... other config

class CustomModel(comfy.model_base.BaseModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)  # ← Fails here
        self.pipeline = {}

# Later:
model = CustomModel(CustomModelConfig(dtype), ...)
model.diffusion_model = my_transformer  # ← Never reached

Debug Logs

AttributeError: 'CustomModel' object has no attribute 'diffusion_model'

Traceback (most recent call last):
  File "/app/ComfyUI/execution.py", line 527, in execute
    output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(...)
  File "/app/ComfyUI/execution.py", line 331, in get_output_data
    return_values = await _async_map_node_over_list(...)
  File "/app/ComfyUI/execution.py", line 305, in _async_map_node_over_list
    await process_inputs(input_dict, i)
  File "/app/ComfyUI/execution.py", line 293, in process_inputs
    result = f(**inputs)
  File ".../custom_node.py", line XX, in __init__
    super().__init__(*args, **kwargs)
  File "/app/ComfyUI/comfy/model_base.py", line 153, in __init__
    comfy.model_management.archive_model_dtypes(self.diffusion_model)
  File "/root/.pyenv/versions/3.11.13/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1965, in __getattr__
    raise AttributeError(
AttributeError: 'CustomModel' object has no attribute 'diffusion_model'

Other

[ivanprado]

It seems the bug was introduced in this commit four days ago

[comfyanonymous]

This should fix it: #12315