VQModel VAE

2023-10-21 22:22:28 +02:00
parent 1c72bd4ae7
commit dcae5c25b0
4 changed files with 234 additions and 1 deletions
@@ -3,7 +3,6 @@ import comfy.sd
 import comfy.utils
 from comfy import model_management
 from comfy import diffusers_convert
-from .models.kl import AutoencoderKL

 vae_dtype_dict = {
 	"auto" : model_management.vae_device(),
@@ -22,7 +21,11 @@ class EXVAE(comfy.sd.VAE):
 		self.latent_scale = model_conf["embed_scale"]

 		if model_conf["type"] == "AutoencoderKL":
+			from .models.kl import AutoencoderKL
 			model = AutoencoderKL(config=model_conf)
+		if model_conf["type"] == "VQModel":
+			from .models.vq import VQModel
+			model = VQModel(config=model_conf)

 		self.first_stage_model = model.eval()
 		m, u = self.first_stage_model.load_state_dict(sd, strict=False)