Skip to content

Commit

Permalink
Fix unknown VAE being detected as the mochi VAE.
Browse files Browse the repository at this point in the history
comfyanonymous committed Nov 5, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 69694f4 commit 8afb97c
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion comfy/sd.py
Original file line number Diff line number Diff line change
@@ -245,7 +245,7 @@ def __init__(self, sd=None, device=None, config=None, dtype=None):
self.process_output = lambda audio: audio
self.process_input = lambda audio: audio
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
elif "blocks.2.blocks.3.stack.5.weight" in sd or "decoder.blocks.2.blocks.3.stack.5.weight" in sd or "layers.4.layers.1.attn_block.attn.qkv.weight" in sd or "encoder.layers.4.layers.1.attn_block.attn.qkv.weight": #genmo mochi vae
elif "blocks.2.blocks.3.stack.5.weight" in sd or "decoder.blocks.2.blocks.3.stack.5.weight" in sd or "layers.4.layers.1.attn_block.attn.qkv.weight" in sd or "encoder.layers.4.layers.1.attn_block.attn.qkv.weight" in sd: #genmo mochi vae
if "blocks.2.blocks.3.stack.5.weight" in sd:
sd = comfy.utils.state_dict_prefix_replace(sd, {"": "decoder."})
if "layers.4.layers.1.attn_block.attn.qkv.weight" in sd:

0 comments on commit 8afb97c

Please sign in to comment.