From 99bff66e3983c64e2390773d06a33175799630db Mon Sep 17 00:00:00 2001 From: Zhaolun Zou Date: Tue, 17 Dec 2024 13:34:10 +0800 Subject: [PATCH] Add key check before dcae key mapping to prevent breaking existing weights; Added both weight addresses to README.md --- README.md | 2 +- VAE/loader.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ab4a667..28cccbe 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ https://github.com/Efficient-Large-Model/ComfyUI_ExtraModels 2. Place them in your checkpoints folder 3. Load them with the correct PixArt checkpoint loader 4. Use the "Gemma Loader" node - it should automatically download the requested model from Huggingface - Recommended to use the 4bit quantized model on CPU when low on memory. -5. Download the VAE from [here](https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors) and place it in your VAE folder after renaming it. +5. Download the VAE from [here](https://huggingface.co/Efficient-Large-Model/Sana_1600M_1024px_diffusers/blob/main/vae/diffusion_pytorch_model.safetensors) or [here](https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors) and place it in your VAE folder after renaming it. 6. Use either the "Empty Sana Latent Image" or "Empty DCAE Latent Image" node for the latent input when doing txt2img. [Sample workflow](https://github.com/user-attachments/files/18027854/SanaV1.json) diff --git a/VAE/loader.py b/VAE/loader.py index d3809ed..874f224 100644 --- a/VAE/loader.py +++ b/VAE/loader.py @@ -34,8 +34,9 @@ class EXVAE(comfy.sd.VAE): model = MoVQ(model_conf) elif model_conf["type"] == "DCAE": from .models.dcae import DCAE - from .models import dcae_key_mapping - sd = dcae_key_mapping.convert_sd(sd) + if 'decoder.project_out.op_list.0.bias' in sd: + from .models import dcae_key_mapping + sd = dcae_key_mapping.convert_sd(sd) model = DCAE(**model_conf) else: raise NotImplementedError(f"Unknown VAE type '{model_conf['type']}'")