misc: add quant and attn backend -> step mask example (#454)

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: support quantize and attn backend for flux example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

* misc: add quant and attn backend -> step mask example

Author: DefTruth

README.md

@@ -11,14 +11,13 @@
     </h2>
   </p>
 
-
-|Baseline|SCM S S*|SCM S D*|SCM F D*|SCM U D*|+TaylorSeer|+compile| 
+|Baseline|SCM S S*|SCM F D*|SCM U D*|+TS|+compile|+FP8*|   
 |:---:|:---:|:---:|:---:|:---:|:---:|:---:|
-|24.85s|15.4s|17.1s|11.4s|8.2s|8.2s|7.1s|
-|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.NONE.png" width=90px>|<img src="assets/steps_mask/static.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.15_SCM1111111101110011100110011000_dynamic_T0O0_S8.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.2_SCM1111110100010000100000100000_dynamic_T0O0_S15.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.3_SCM111101000010000010000001000000_dynamic_T0O0_S19.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=90px>|
+|24.85s|15.4s|11.4s|8.2s|8.2s|7.1s|4.5s|
+|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.NONE.png" width=90px>|<img src="assets/steps_mask/static.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.2_SCM1111110100010000100000100000_dynamic_T0O0_S15.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.3_SCM111101000010000010000001000000_dynamic_T0O0_S19.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=90px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=90px>|<img src="./assets/steps_mask/flux.C1_Q1_float8_DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=90px>|
 
 <p align="center">
-  Scheme: <b>DBCache + SCM(steps_computation_mask) + TaylorSeer</b>, L20x1, S*: static cache, <b>D*: dynamic cache</b>, <br><b>S</b>: Slow, <b>F</b>: Fast, <b>U</b>: Ultra Fast, FLUX.1-Dev, Steps: 28, Prompt: "A cat holding a sign that says hello world"
+  Scheme: <b>DBCache + SCM(steps_computation_mask) + TS(TaylorSeer) + FP8*</b>, L20x1, S*: static cache, <br><b>D*: dynamic cache</b>, <b>S</b>: Slow, <b>F</b>: Fast, <b>U</b>: Ultra Fast, <b>TS</b>: TaylorSeer, <b>FP8*</b>: FP8 DQ + Sage, <b>FLUX.1</b>-Dev
 </p>
 
 <img src=https://github.com/vipshop/cache-dit/raw/main/assets/speedup_v4.png>

assets/steps_mask/flux.C1_Q1_float8_DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png

@@ -615,7 +615,7 @@ cache_dit.enable_cache(
 
 <div align="center">
 
-|Baseline(L20x1)|F1B0 (0.12)|+TaylorSeer|F1B0 (0.15)|+TaylorSeer|+compile| 
+|Baseline(L20x1)|F1B0 (0.12)|+TaylorSeer|F1B0 (0.15)|+TaylorSeer|+compile|  
 |:---:|:---:|:---:|:---:|:---:|:---:|
 |24.85s|12.85s|12.86s|10.27s|10.28s|8.48s|
 |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=140px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.12_S14_T12.85s.png width=140px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.12_S14_T12.86s.png width=140px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.15_S17_T10.27s.png width=140px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T10.28s.png width=140px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T8.48s.png width=140px>|
@@ -674,13 +674,13 @@ As we can observe, in the case of **static cache**, the image of `SCM Slow S*` (
 
 <div align="center">
 
-|Baseline(L20x1)|SCM Slow S*|SCM Slow D*|SCM Fast D*|SCM Ultra D*|+TaylorSeer|+compile| 
-|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
-|24.85s|15.4s|17.1s|11.4s|8.2s|8.2s|7.1s|
-|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.NONE.png" width=110px>|<img src="../assets/steps_mask/static.png" width=110px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.15_SCM1111111101110011100110011000_dynamic_T0O0_S8.png" width=110px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.2_SCM1111110100010000100000100000_dynamic_T0O0_S15.png" width=110px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.3_SCM111101000010000010000001000000_dynamic_T0O0_S19.png" width=110px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=110px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=110px>|
+|Baseline|SCM S S*|SCM S D*|SCM F D*|SCM U D*|+TS|+compile|+FP8 +Sage|  
+|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+|24.85s|15.4s|17.1s|11.4s|8.2s|8.2s|7.1s|4.5s|
+|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.NONE.png" width=95px>|<img src="../assets/steps_mask/static.png" width=95px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.15_SCM1111111101110011100110011000_dynamic_T0O0_S8.png" width=95px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.2_SCM1111110100010000100000100000_dynamic_T0O0_S15.png" width=95px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.3_SCM111101000010000010000001000000_dynamic_T0O0_S19.png" width=95px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=95px>|<img src="https://github.com/vipshop/cache-dit/raw/main/assets/steps_mask/flux.DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=95px>|<img src="../assets/steps_mask/flux.C1_Q1_float8_DBCache_F1B0_W8I1M0MC0_R0.35_SCM111101000010000010000001000000_dynamic_T1O1_S19.png" width=95px>|
 
 <p align="center">
-    DBCache + SCM(steps_computation_mask) + TaylorSeer, <b> L20x1 </b>, <br>S*: static cache, D*: dynamic cache, Steps: 28, "A cat holding a sign that says hello world"
+  Scheme: <b>DBCache + SCM(steps_computation_mask) + TaylorSeer</b>, L20x1, S*: static cache, <b>D*: dynamic cache</b>, <br><b>S</b>: Slow, <b>F</b>: Fast, <b>U</b>: Ultra Fast, <b>TS</b>: TaylorSeer, FP8: FP8 DQ, Sage: SageAttention, <b>FLUX.1-Dev</b>, <br>Steps: 28, HxW=1024x1024, Prompt: "A cat holding a sign that says hello world"
 </p>
 
 |DBCache + SCM Slow S*|DBCache + SCM Ultra D* + TaylorSeer + compile| 

docs/User_Guide.md

@@ -5,8 +5,8 @@
 
 import time
 import torch
-from diffusers import FluxPipeline
-from utils import get_args
+from diffusers import FluxPipeline, FluxTransformer2DModel
+from utils import get_args, strify
 import cache_dit
 
 
@@ -62,7 +62,7 @@
         "black-forest-labs/FLUX.1-dev",
     ),
     torch_dtype=torch.bfloat16,
-).to("cuda")
+)
 
 if args.cache:
     from cache_dit import DBCacheConfig, TaylorSeerCalibratorConfig
@@ -94,10 +94,36 @@
         ),
     )
 
+assert isinstance(pipe.transformer, FluxTransformer2DModel)
+if args.quantize:
+    pipe.transformer = cache_dit.quantize(
+        pipe.transformer,
+        quant_type=args.quantize_type,
+        exclude_layers=[
+            "embedder",
+            "embed",
+        ],
+    )
+    pipe.text_encoder_2 = cache_dit.quantize(
+        pipe.text_encoder_2,
+        quant_type=args.quantize_type,
+    )
+    print(f"Applied quantization: {args.quantize_type} to Transformer and Text Encoder 2.")
+
+pipe.to("cuda")
+
+if args.attn is not None:
+    if hasattr(pipe.transformer, "set_attention_backend"):
+        pipe.transformer.set_attention_backend(args.attn)
+        print(f"Set attention backend to {args.attn}")
+
 
 if args.compile:
     cache_dit.set_compile_configs()
     pipe.transformer = torch.compile(pipe.transformer)
+    pipe.text_encoder = torch.compile(pipe.text_encoder)
+    pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2)
+    pipe.vae = torch.compile(pipe.vae)
 
 
 def run_pipe():
@@ -121,7 +147,7 @@ def run_pipe():
 cache_dit.summary(pipe)
 
 time_cost = end - start
-save_path = f"flux.{cache_dit.strify(pipe)}.png"
+save_path = f"flux.{strify(args, pipe)}.png"
 print(f"Time cost: {time_cost:.2f}s")
 print(f"Saving image to {save_path}")
 image.save(save_path)
@@ -134,3 +160,4 @@ def run_pipe():
 # python3 run_steps_mask.py --cache --Fn 1 --step-mask u --step-policy dynamic --rdt 0.30
 # python3 run_steps_mask.py --cache --Fn 1 --step-mask u --step-policy dynamic --rdt 0.30 --taylorseer --taylorseer-order 1
 # python3 run_steps_mask.py --cache --Fn 1 --step-mask u --step-policy dynamic --rdt 0.30 --compile --taylorseer --taylorseer-order 1
+# python3 run_steps_mask.py --cache --Fn 1 --step-mask u --step-policy dynamic --rdt 0.35 --compile --taylorseer --taylorseer-order 1 --quantize --quantize-type float8 --attn sage

examples/api/run_steps_mask.py

@@ -48,6 +48,13 @@
         pipe.transformer.set_attention_backend(args.attn)
         print(f"Set attention backend to {args.attn}")
 
+if args.compile:
+    cache_dit.set_compile_configs()
+    pipe.transformer = torch.compile(pipe.transformer)
+    pipe.text_encoder = torch.compile(pipe.text_encoder)
+    pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2)
+    pipe.vae = torch.compile(pipe.vae)
+
 
 def run_pipe():
     image = pipe(
@@ -59,14 +66,6 @@ def run_pipe():
     ).images[0]
     return image
 
-
-if args.compile:
-    cache_dit.set_compile_configs()
-    pipe.transformer = torch.compile(pipe.transformer)
-    pipe.text_encoder = torch.compile(pipe.text_encoder)
-    pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2)
-    pipe.vae = torch.compile(pipe.vae)
-
 # warmup
 _ = run_pipe()