Update src/pipeline.py

src/pipeline.py

@@ -38,10 +38,8 @@ torch.backends.cudnn.enabled = True
 
 # globals
 Pipeline = None
-ckpt_id = "black-forest-labs/FLUX.1-schnell"
-ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
-TinyVAE = "madebyollin/taef1"
-TinyVAE_REV = "2d552378e58c9c94201075708d7de4e1163b2689"
+ckpt_id = "manbeast3b/Flux.1.schnell-quant2"
+ckpt_revision = "44eb293715147878512da10bf3bc47cd14ec8c55" 
 
 def empty_cache():
     gc.collect()
@@ -50,6 +48,8 @@ def empty_cache():
     torch.cuda.reset_peak_memory_stats()
 
 def load_pipeline() -> Pipeline:
+    vae = AutoencoderKL.from_pretrained(ckpt_id,revision=ckpt_revision, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16,)
+    quantize_(vae, int8_weight_only())
     text_encoder_2 = T5EncoderModel.from_pretrained("manbeast3b/flux.1-schnell-full1", revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146", subfolder="text_encoder_2",torch_dtype=torch.bfloat16)
     path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--flux.1-schnell-full1/snapshots/cb1b599b0d712b9aab2c4df3ad27b050a27ec146/transformer")
     transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)