change vae

app_onediff.py

@@ -8,11 +8,9 @@ import json
 from diffusers import (
     DDPMScheduler,
     AutoPipelineForText2Image,
-    AutoencoderTiny,
+    AutoencoderKL,
 )
 
-import sys
-
 os.system("python3 -m pip --no-cache-dir install --pre nexfort -f https://github.com/siliconflow/nexfort_releases/releases/expanded_assets/torch2.4.1_cu121")
 os.system("git clone https://github.com/siliconflow/onediff.git")
 os.system("cd onediff && python3 -m pip install .")
@@ -33,11 +31,10 @@ def nexfort_compile(torch_module: torch.nn.Module):
 BASE_MODEL = "stabilityai/sdxl-turbo"
 device = "cuda"
 
-vae = AutoencoderTiny.from_pretrained(
-    'madebyollin/taesdxl',
-    use_safetensors=True,
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", 
     torch_dtype=torch.float16,
-).to('cuda')
+)
 base_pipe = AutoPipelineForText2Image.from_pretrained(
     BASE_MODEL, 
     vae=vae,