Update src/pipeline.py

src/pipeline.py

@@ -4,7 +4,7 @@ from typing import TypeAlias
 
 import torch
 from PIL.Image import Image
-from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL
+from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, AutoencoderTiny
 from huggingface_hub.constants import HF_HUB_CACHE
 from pipelines.models import TextToImageRequest
 from torch import Generator
@@ -23,7 +23,7 @@ def load_pipeline() -> Pipeline:
 
     path1 = os.path.join(HF_HUB_CACHE, "models--aifeifei798--taef1/snapshots/319774e8256cc42d98b3c2790ac2c5edc8a73389")
 
-    vae = AutoencoderKL.from_pretrained(path1, torch_dtype=torch.bfloat16,)
+    vae = AutoencoderTiny.from_pretrained(path1, torch_dtype=torch.bfloat16,)
 
     path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")