Update src/pipeline.py

src/pipeline.py

@@ -12,7 +12,7 @@ from pipelines.models import TextToImageRequest
 from torch import Generator
 import time
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
-# from torchao.quantization import quantize_, int8_weight_only
+from torchao.quantization import quantize_, int8_weight_only
 import torch.nn.utils.prune as prune
 Pipeline = None
 
@@ -31,15 +31,16 @@ def load_pipeline() -> Pipeline:
         ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16
     )#.to("cpu")
     quantize_(vae, int8_weight_only())
+    for name, module in pipeline.vae.named_modules():
+      if isinstance(module, torch.nn.Linear):
+          prune.l1_unstructured(module, 'weight', amount=0.1)
     pipeline = DiffusionPipeline.from_pretrained(
         ckpt_id,
         vae=vae,
         torch_dtype=dtype,
         )#.to("cpu")
 
-    # for name, module in pipeline.transformer.named_modules():
-    #   if isinstance(module, torch.nn.Linear):
-    #       prune.l1_unstructured(module, 'weight', amount=0.05)
+
     
     # pipeline.to(device)
     pipeline.enable_sequential_cpu_offload()