Update src/pipeline.py

src/pipeline.py

@@ -1,4 +1,6 @@
 from torch import Generator
+from diffusers.image_processor import VaeImageProcessor
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 import torch
 from PIL.Image import Image
 from pipelines.models import TextToImageRequest
@@ -9,8 +11,10 @@ import os
 from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
 import torch._dynamo
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 HOME = os.environ["HOME"]
+
 Pipeline = None
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 
@@ -22,10 +26,9 @@ def empty_cache():
     
 def load_pipeline() -> Pipeline:
     empty_cache()
-    dtype, device = torch.bfloat16, "cuda"
     text_encoder = CLIPTextModel.from_pretrained(ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
     quantize_(text_encoder, int8_weight_only())
-    vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_e3m2", torch_dtype=torch.bfloat16)
+    vae = AutoencoderTiny.from_pretrained("aifeifei798/taef1", torch_dtype=torch.bfloat16)
     quantize_(vae, int8_weight_only())
     text_encoder_2 = T5EncoderModel.from_pretrained("city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16)
     quantize_(text_encoder_2, int8_weight_only())