Upload src/pipeline.py with huggingface_hub

src/pipeline.py

@@ -28,20 +28,20 @@ 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)
+    
     pipeline = FluxPipeline.from_pretrained(
         CHECKPOINT,
         revision=REVISION,
         transformer=transformer,
         local_files_only=True,
         torch_dtype=torch.bfloat16,
-    ).to("cuda")  
+    )
 
-#    pipeline.vae.to(memory_format=torch.channels_last) 
+    pipeline.to(memory_format=torch.channels_last)
     quantize_(pipeline.vae, int8_weight_only())
-#    pipeline.vae.compile()
     pipeline.to("cuda")
-    for _ in range(2):
+    for _ in range(4):
         pipeline("cat", num_inference_steps=4)
 
     return pipeline