Fix "Load Checkpoint with config" node.

comfyanonymous · Aug 30, 2023 · fe4c074 · fe4c074 · HellaRelevant · Aug 30, 2023
1 parent d70b0bc
commit fe4c074
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/comfy/sd.py b/comfy/sd.py
@@ -22,6 +22,7 @@
 import comfy.model_patcher
 import comfy.lora
 import comfy.t2i_adapter.adapter
+import comfy.supported_models_base
 
 def load_model_weights(model, sd):
     m, u = model.load_state_dict(sd, strict=False)
@@ -348,10 +349,11 @@ class WeightsLoader(torch.nn.Module):
     class EmptyClass:
         pass
 
-    model_config = EmptyClass()
-    model_config.unet_config = unet_config
+    model_config = comfy.supported_models_base.BASE({})
+
     from . import latent_formats
     model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)
+    model_config.unet_config = unet_config
 
     if config['model']["target"].endswith("LatentInpaintDiffusion"):
         model = model_base.SDInpaint(model_config, model_type=model_type)

diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py
@@ -1,6 +1,7 @@
 import torch
 from . import model_base
 from . import utils
+from . import latent_formats
 
 
 def state_dict_key_replace(state_dict, keys_to_replace):
@@ -34,6 +35,7 @@ class BASE:
     clip_vision_prefix = None
     noise_aug_config = None
     beta_schedule = "linear"
+    latent_format = latent_formats.LatentFormat
 
     @classmethod
     def matches(s, unet_config):