From da205c0a2a2d1d5779b2c5273e70481e888d0ef1 Mon Sep 17 00:00:00 2001 From: Ruslan Rakhimov Date: Wed, 2 Sep 2026 12:27:10 +0300 Subject: [PATCH] Fix two unreachable keys in VRAM_MANAGEMENT_MODULE_MAPS The map is looked up by the model_class string registered in model_configs.py, so an entry keyed on any other name is never selected and the model silently falls back to whole-model wrapping. - joyai_image_dit.Transformer3DModel: no such class exists; the registered class is JoyAIImageDiT. - qwen_image_controlnet.BlockWiseControlBlock: that class is the inner block; the registered class is QwenImageBlockWiseControlNet. --- diffsynth/configs/vram_management_module_maps.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index b7636771..aa757877 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -80,7 +80,7 @@ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.qwen_image_vae.QwenImageRMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.qwen_image_controlnet.BlockWiseControlBlock": { + "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet": { "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", }, @@ -339,7 +339,7 @@ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", "transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", }, - "diffsynth.models.joyai_image_dit.Transformer3DModel": { + "diffsynth.models.joyai_image_dit.JoyAIImageDiT": { "diffsynth.models.joyai_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", "diffsynth.models.joyai_image_dit.ModulateWan": "diffsynth.core.vram.layers.AutoWrappedModule", "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",