Update src/pipeline.py

src/pipeline.py

@@ -44,7 +44,6 @@ ckpt_revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146"
 def load_pipeline() -> Pipeline:
     # model_name = "manbeast3b/flux.1-schnell-full1"
     # text_enc_revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146"
-
     # text_encoder_2 = T5EncoderModel.from_pretrained(
     #     model_name,
     #     revision=text_enc_revision,
@@ -83,7 +82,7 @@ def load_pipeline() -> Pipeline:
     ).to("cuda")
     # pipeline.vae = torch.compile(vae)
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
-    quantize_(pipeline.vae, float8_weight_only())
+    
     pipeline.to(memory_format=torch.channels_last)
     
     warmup_ = "controllable varied focus thai warriors entertainment blue golden pink soft tough padthai"
@@ -96,6 +95,7 @@ def load_pipeline() -> Pipeline:
             num_inference_steps=4,
             max_sequence_length=256
         )
+    quantize_(pipeline.vae, float8_weight_only())
     # pipeline("")
     return pipeline