Further optimize + quantize

app.py

@@ -15,6 +15,8 @@ import gradio as gr
 import spaces
 import torch
 from diffusers import FluxPipeline
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 
 from zerogpu import aoti_compile
 
@@ -26,6 +28,12 @@ print('FluxPipeline.from_pretrained', -(t0 - (t0 := datetime.now())))
 @spaces.GPU(duration=1500)
 def compile_transformer():
 
+    pipeline.transformer.fuse_qkv_projections()
+    pipeline.vae.fuse_qkv_projections()
+    pipeline.vae.to(memory_format=torch.channels_last)
+
+    quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
+
     def _example_tensor(*shape):
         return torch.randn(*shape, device='cuda', dtype=torch.bfloat16)
 

requirements.txt

@@ -3,3 +3,4 @@ diffusers
 transformers
 sentencepiece
 protobuf
+torchao