Update src/pipeline.py

src/pipeline.py

@@ -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 vae.named_modules():
-      if isinstance(module, torch.nn.Linear):
-          prune.random_unstructured(module, name="weight", amount=0.2)
+    
           # 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.random_unstructured(module, name="weight", amount=0.2)