Update inference_manager.py

inference_manager.py

@@ -151,9 +151,13 @@ class InferenceManager:
             vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
             pipe = StableDiffusionPipeline.from_pretrained(ckpt_dir, vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True)
         else:
-            
-            #vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
-            vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.bfloat16)
+            use_vae = cfg.get("vae", "")
+            if not use_vae:
+                vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.bfloat16)
+            elif use_vae == "tae":
+                vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.bfloat16)
+            else:
+                vae = AutoencoderTiny.from_pretrained(use_vae, torch_dtype=torch.bfloat16)
             print(ckpt_dir)
             pipe = DiffusionPipeline.from_pretrained(
                 ckpt_dir,