flux.2

2025-11-27 00:10:55 +08:00
parent 0eead33ed7
commit 0b527c460f
11 changed files with 2228 additions and 1 deletions
--- a/diffsynth/configs/vram_management_module_maps.py
+++ b/diffsynth/configs/vram_management_module_maps.py
@@ -150,4 +150,9 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
        "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
    },
+    "diffsynth.models.flux2_dit.Flux2DiT": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
 }