Update src/pipeline.py

src/pipeline.py

@@ -1,11 +1,11 @@
-#7
+#7.1
 from huggingface_hub.constants import HF_HUB_CACHE
 from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
 import torch
 import torch._dynamo
 import gc
 import os
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+from diffusers import FluxPipeline, AutoencoderTiny
 from PIL.Image import Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
@@ -27,6 +27,7 @@ def load_pipeline() -> Pipeline:
     pipeline.to("cuda")
     quantize_(pipeline.vae, int8_weight_only())
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
+    pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune", fullgraph=True, dynamic=True)
     for _ in range(3):
         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