Fix checkpoint loader with config.

ltdrdata · Oct 28, 2023 · 2a134bf · 2a134bf
1 parent e60ca69
commit 2a134bf
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/comfy/sd.py b/comfy/sd.py
@@ -388,11 +388,13 @@ class EmptyClass:
         if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
             clip_target.clip = sd2_clip.SD2ClipModel
             clip_target.tokenizer = sd2_clip.SD2Tokenizer
+            clip = CLIP(clip_target, embedding_directory=embedding_directory)
+            w.cond_stage_model = clip.cond_stage_model.clip_h
         elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
             clip_target.clip = sd1_clip.SD1ClipModel
             clip_target.tokenizer = sd1_clip.SD1Tokenizer
-        clip = CLIP(clip_target, embedding_directory=embedding_directory)
-        w.cond_stage_model = clip.cond_stage_model
+            clip = CLIP(clip_target, embedding_directory=embedding_directory)
+            w.cond_stage_model = clip.cond_stage_model.clip_l
         load_clip_weights(w, state_dict)
 
     return (comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)

diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
@@ -472,11 +472,11 @@ def untokenize(self, token_weight_pair):
 
 
 class SD1ClipModel(torch.nn.Module):
-    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel):
+    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
         super().__init__()
         self.clip_name = clip_name
         self.clip = "clip_{}".format(self.clip_name)
-        setattr(self, self.clip, clip_model(device=device, dtype=dtype))
+        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
 
     def clip_layer(self, layer_idx):
         getattr(self, self.clip).clip_layer(layer_idx)

diff --git a/comfy/sd2_clip.py b/comfy/sd2_clip.py
@@ -21,5 +21,5 @@ def __init__(self, embedding_directory=None):
         super().__init__(embedding_directory=embedding_directory, clip_name="h", tokenizer=SD2ClipHTokenizer)
 
 class SD2ClipModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None):
-        super().__init__(device=device, dtype=dtype, clip_name="h", clip_model=SD2ClipHModel)
+    def __init__(self, device="cpu", dtype=None, **kwargs):
+        super().__init__(device=device, dtype=dtype, clip_name="h", clip_model=SD2ClipHModel, **kwargs)