fix alias rev map for quant types (#629)

* chore: fix alias for quant types

* chore: fix alias for quant types

* chore: fix alias for quant types

* chore: fix alias for quant types

Author: DefTruth

README.md

@@ -76,13 +76,13 @@ You can install the stable release of cache-dit from PyPI, or the latest develop
 ## 🔥Supported DiTs
 
 > [!Tip]   
-> One Model Series may contain many pipelines. cache-dit applies optimizations at the Transformer level; thus, any pipelines that include the supported transformer are already supported by cache-dit. ✅: supported now; ✖️: not supported now; **[C-P](./)**: Context Parallelism; **[T-P](./)**: Tensor Parallelism; **[TE-P](./)**: Text Encoder Parallelism; **[CN-P](./)**: ControlNet Parallelism;  **[VAE-P](./)**: VAE Parallelism (TODO).
+> One Model Series may contain many pipelines. cache-dit applies optimizations at the Transformer level; thus, any pipelines that include the supported transformer are already supported by cache-dit. ✅: supported now; ✖️: not supported now; **[🤖Q](https://github.com/nunchaku-tech/nunchaku)**: **[nunchaku](https://github.com/nunchaku-tech/nunchaku)** w/ SVDQ W4A4; **[C-P](./)**: Context Parallelism; **[T-P](./)**: Tensor Parallelism; **[TE-P](./)**: Text Encoder Parallelism; **[CN-P](./)**: ControlNet Parallelism;  **[VAE-P](./)**: VAE Parallelism (TODO).
 
 <div align="center">
 
 | 📚Supported DiTs: `🤗65+` | Cache  | C-P | T-P | TE-P | CN-P | VAE-P |
 |:---:|:---:|:---:|:---:|:---:|:---:|:---:|
-| Z-Image-Turbo `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Z-Image-Turbo `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
 | Qwen-Image-Layered | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |
 | Qwen-Image-Edit-2511-Lightning | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |
 | Qwen-Image-Edit-2511 | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |
@@ -114,14 +114,15 @@ You can install the stable release of cache-dit from PyPI, or the latest develop
 | HunyuanImage-2.1 | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |
 | HunyuanVideo-1.5 | ✅ | ✖️ | ✖️ | ✅ | ✖️ | ✖️ |
 | HunyuanVideo | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |
-| FLUX.1-dev `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| FLUX.1-Fill-dev `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen-Image `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen-Image-Edit `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen-Image-Edit-2509 `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen-Image-Lightning `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen...Edit-Lightning `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
-| Qwen...Edit-2509-Lightning `⚡️Nunchaku` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| FLUX.1-dev `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| FLUX.1-Fill-dev `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| FLUX.1-Kontext-dev `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image-Edit `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image-Edit-2509 `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image-Lightning `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image-Edit-Lightning `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
+| Qwen-Image-Edit-2509-Lightning `🤖Q` | ✅ | ✅ | ✖️ | ✅ | ✖️ | ✖️ |
 | SkyReels-V2-T2V | ✅ | ✅  | ✅  | ✅ | ✖️ | ✖️ |
 | LongCat-Video | ✅ | ✖️ | ✖️ | ✅ | ✖️ | ✖️ |
 | ChronoEdit-14B | ✅ | ✅ | ✅ | ✅ | ✖️ | ✖️ |

src/cache_dit/quantize/torchao/quantize_ao.py

@@ -1,4 +1,5 @@
 import torch
+import copy
 from typing import Callable, Optional, List
 from cache_dit.utils import maybe_empty_cache
 from cache_dit.logger import init_logger
@@ -35,6 +36,13 @@ def quantize_ao(
         "int4_weight_only": "int4_w4a16_wo",
         "int4_wo": "int4_w4a16_wo",
     }
+    alias_map_rev = copy.deepcopy(alias_map)
+    # remove duplicates *_wo in rev map
+    for key in list(alias_map_rev.keys()):
+        if key.endswith("_wo"):
+            alias_map_rev.pop(key)
+    alias_map_rev = {v: k for k, v in alias_map_rev.items()}
+
     if quant_type.lower() in alias_map:
         quant_type = alias_map[quant_type.lower()]
 
@@ -187,7 +195,6 @@ def _quant_config():
 
     maybe_empty_cache()
 
-    alias_map_rev = {v: k for k, v in alias_map.items()}
     if quant_type in alias_map_rev:
         quant_type = alias_map_rev[quant_type]