Refactor model loading to use AutoencoderKL for VAE

app.py

@@ -58,7 +58,7 @@ def load_sd():
         except:
             model["pipeline"] = AutoPipelineForText2Image.from_pretrained(
                 model['repo_id'],
-                vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device),
+                vae = AutoencoderKL.from_pretrained("AutoencoderKL", subfolder="vae", torch_dtype=torch.bfloat16).to(device),
                 torch_dtype = model['compute_type'],
                 safety_checker = None
             ).to(device)
@@ -67,7 +67,7 @@ def load_sd():
 
 
     # VAE n Refiner
-    flux_vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device)
+    flux_vae = AutoencoderKL.from_pretrained("AutoencoderKL", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
     sdxl_vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device)
     refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=sdxl_vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to(device)
     refiner.enable_model_cpu_offload()