Update src/pipeline.py

src/pipeline.py

@@ -40,8 +40,6 @@ torch.backends.cudnn.enabled = True
 Pipeline = None
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
-TinyVAE = "madebyollin/taef1"
-TinyVAE_REV = "2d552378e58c9c94201075708d7de4e1163b2689"
 
 def empty_cache():
     gc.collect()
@@ -50,13 +48,13 @@ def empty_cache():
     torch.cuda.reset_peak_memory_stats()
 
 def load_pipeline() -> Pipeline:
-    text_encoder_2 = T5EncoderModel.from_pretrained("manbeast3b/flux.1-schnell-full1", revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146", subfolder="text_encoder_2",torch_dtype=torch.bfloat16)
+    text_encoder_2 = T5EncoderModel.from_pretrained("manbeast3b/flux.1-schnell-full1", revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146", subfolder="text_encoder_2",torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
     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)
+    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False).to(memory_format=torch.channels_last)
+    quantize_(AutoencoderKL.from_pretrained(ids,revision=Revision, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16,), int8_weight_only())
     pipeline = FluxPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
     pipeline.to("cuda")
-    pipeline.to(memory_format=torch.channels_last)
-    for _ in range(1):
+    with torch.inference_mode():
         pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
     return pipeline