Update adapter_model.py

adapter_model.py

@@ -1,6 +1,6 @@
 from dotenv import load_dotenv
 import os 
-from diffusers import StableDiffusionInpaintPipelineLegacy,StableDiffusionControlNetPipeline,ControlNetModel, DDIMScheduler,AutoencoderKL
+from diffusers import StableDiffusionInpaintPipelineLegacy, StableDiffusionControlNetPipeline, ControlNetModel, DDIMScheduler, AutoencoderKL, ConsistencyDecoderVAE
 import torch
 from ip_adapter import IPAdapter
 
@@ -21,9 +21,9 @@ IP_CKPT_PATH = str(os.getenv(
 DEVICE = str(os.getenv(
     "DEVICE")
     )
-
-from diffusers import ConsistencyDecoderVAE
-
+OPENAI_CONSISTENCY_VAE = str(os.getenv(
+    "OPENAI_CONSISTENCY_VAE")
+    )
 
 
 noise_scheduler = DDIMScheduler(
@@ -35,8 +35,12 @@ noise_scheduler = DDIMScheduler(
     set_alpha_to_one=False,
     steps_offset=1,
 )
-vae = AutoencoderKL.from_pretrained(VAE_MODEL_PATH).to(dtype=torch.float16)
-vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
+
+if OPENAI_CONSISTENCY_VAE == "ENABLE":
+    vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
+
+else :
+    vae = AutoencoderKL.from_pretrained(VAE_MODEL_PATH).to(dtype=torch.float16)
 
 
 class MODEL: