Upload src/pipeline.py with huggingface_hub

src/pipeline.py

@@ -28,20 +28,27 @@ def load_pipeline() -> Pipeline:
             path, 
             use_safetensors=False,
             local_files_only=True,
-            torch_dtype=torch.bfloat16)
+            torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
+    vae = AutoencoderTiny.from_pretrained(
+            TinyVAE,
+            revision=TinyVAE_REV,
+            local_files_only=True,
+            torch_dtype=torch.bfloat16
+            )    
     pipeline = FluxPipeline.from_pretrained(
         CHECKPOINT,
         revision=REVISION,
         transformer=transformer,
+        vae=vae,
         local_files_only=True,
         torch_dtype=torch.bfloat16,
-    ).to("cuda")  
+    )
 
-    pipeline.vae.to(memory_format=torch.channels_last) 
-#    quantize_(pipeline.vae, int8_weight_only())
-#    pipeline.vae.compile()
+    pipeline.to(memory_format=torch.channels_last)
+    pipeline.enable_vae_slicing()
     pipeline.to("cuda")
-    for _ in range(2):
+    # quantize_(pipeline.vae, int8_weight_only())
+    for _ in range(4):
         pipeline("cat", num_inference_steps=4)
 
     return pipeline