diff --git a/comfy/model_base.py b/comfy/model_base.py index dcbf12074..3bb54f59e 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -147,11 +147,11 @@ class BaseModel(torch.nn.Module): self.diffusion_model.to(memory_format=torch.channels_last) logging.debug("using channels last mode for diffusion model") logging.info("model weight dtype {}, manual cast: {}".format(self.get_dtype(), self.manual_cast_dtype)) + comfy.model_management.archive_model_dtypes(self.diffusion_model) + self.model_type = model_type self.model_sampling = model_sampling(model_config, model_type) - comfy.model_management.archive_model_dtypes(self.diffusion_model) - self.adm_channels = unet_config.get("adm_in_channels", None) if self.adm_channels is None: self.adm_channels = 0