Add consistency decoder

modules/sd_vae_consistency.py

@@ -0,0 +1,35 @@
+"""
+Consistency Decoder
+Improved decoding for stable diffusion vaes.
+
+https://github.com/openai/consistencydecoder
+"""
+import os
+import torch
+import torch.nn as nn
+
+from modules import devices, paths_internal, shared
+from consistencydecoder import ConsistencyDecoder
+
+
+sd_vae_consistency_models = None
+model_path = os.path.join(paths_internal.models_path, 'consistencydecoder')
+
+
+def decoder_model():
+    global sd_vae_consistency_models
+    if getattr(shared.sd_model, 'is_sdxl', False):
+        raise NotImplementedError("SDXL is not supported for consistency decoder")
+    if sd_vae_consistency_models is not None:
+        sd_vae_consistency_models.ckpt.to(devices.device)
+        return sd_vae_consistency_models
+
+    loaded_model = ConsistencyDecoder(devices.device, model_path)
+    sd_vae_consistency_models = loaded_model
+    return loaded_model
+
+
+def unload():
+    global sd_vae_consistency_models
+    if sd_vae_consistency_models is not None:
+        sd_vae_consistency_models.ckpt.to('cpu')