Merge pull request #103 from zergzzlun/fix/vae-weight-mapping

Key Mapping for DCAE Weight Compatibility

Gemma/nodes.py

@@ -33,7 +33,7 @@ class GemmaLoader:
             devices.append(f"cuda:{k}")
         return {
             "required": {
-                "model_name": (["google/gemma-2-2b-it", "unsloth/gemma-2-2b-it-bnb-4bit"],),
+                "model_name": (["Efficient-Large-Model/gemma-2-2b-it", "google/gemma-2-2b-it", "unsloth/gemma-2-2b-it-bnb-4bit"],),
                 "device": (devices, {"default":"cpu"}),
                 "dtype": (dtypes,),
             }
@@ -56,6 +56,10 @@ class GemmaLoader:
             text_encoder_dir = os.path.join(folder_paths.models_dir, 'text_encoders', 'models--unsloth--gemma-2-2b-it-bnb-4bit')
             if not os.path.exists(os.path.join(text_encoder_dir, 'model.safetensors')):
                 snapshot_download('unsloth/gemma-2-2b-it-bnb-4bit', local_dir=text_encoder_dir)
+        elif model_name == 'Efficient-Large-Model/gemma-2-2b-it':
+            text_encoder_dir = os.path.join(folder_paths.models_dir, 'text_encoders', 'models--Efficient-Large-Model--gemma-2-2b-it')
+            if not os.path.exists(os.path.join(text_encoder_dir, 'model.safetensors')):
+                snapshot_download('Efficient-Large-Model/gemma-2-2b-it', local_dir=text_encoder_dir)
         else:
             raise ValueError('Not implemented!')
 

README.md

@@ -23,7 +23,6 @@ Clone the repository to your custom nodes folder, assuming haven't installed in
 To install the requirements on windows, run these commands in the same window:
 ```
 .\python_embeded\python.exe -s -m pip install -r .\ComfyUI\custom_nodes\ComfyUI_ExtraModels\requirements.txt
-.\python_embeded\python.exe -s -m pip install bitsandbytes --prefer-binary --extra-index-url=https://jllllll.github.io/bitsandbytes-windows-webui
 ```
 
 To update, open the command line window like before and run the following commands:
@@ -40,6 +39,16 @@ Alternatively, use the manager, assuming it has an update function.
 
 [Original Repo](https://github.com/NVlabs/Sana)
 
+> [!CAUTION]
+> As many people have had issues with Sana, it's for now recommended to try the fork by the Sana devs, which auto downloads all models:
+> 
+> [Readme](https://github.com/NVlabs/Sana/blob/main/asset/docs/ComfyUI/comfyui.md) | [Fork repo](https://github.com/Efficient-Large-Model/ComfyUI_ExtraModels)
+
+A full rewrite to have better integration is in progress in [this PR](https://github.com/city96/ComfyUI_ExtraModels/pull/92) but isn't ready yet.
+
+https://github.com/NVlabs/Sana/blob/main/asset/docs/ComfyUI/comfyui.md
+https://github.com/Efficient-Large-Model/ComfyUI_ExtraModels
+
 ### Model info / implementation
 - Uses Gemma2 2B as the text encoder
 - Multiple resolutions and models available
@@ -50,7 +59,7 @@ Alternatively, use the manager, assuming it has an update function.
 2. Place them in your checkpoints folder
 3. Load them with the correct PixArt checkpoint loader
 4. Use the "Gemma Loader" node - it should automatically download the requested model from Huggingface - Recommended to use the 4bit quantized model on CPU when low on memory.
-5. Download the VAE from [here](https://huggingface.co/Efficient-Large-Model/Sana_1600M_1024px_diffusers/blob/main/vae/diffusion_pytorch_model.safetensors) and place it in your VAE folder after renaming it.
+5. Download the VAE from [here](https://huggingface.co/Efficient-Large-Model/Sana_1600M_1024px_diffusers/blob/main/vae/diffusion_pytorch_model.safetensors) or [here](https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors) and place it in your VAE folder after renaming it.
 6. Use either the "Empty Sana Latent Image" or "Empty DCAE Latent Image" node for the latent input when doing txt2img.
 
 [Sample workflow](https://github.com/user-attachments/files/18027854/SanaV1.json)

VAE/loader.py

@@ -34,6 +34,9 @@ class EXVAE(comfy.sd.VAE):
 			model = MoVQ(model_conf)
 		elif model_conf["type"] == "DCAE":
 			from .models.dcae import DCAE
+			if 'decoder.project_out.op_list.0.bias' in sd:
+				from .models import dcae_key_mapping
+				sd = dcae_key_mapping.convert_sd(sd)
 			model = DCAE(**model_conf)
 		else:
 			raise NotImplementedError(f"Unknown VAE type '{model_conf['type']}'")

VAE/models/dcae_key_mapping.py

@@ -0,0 +1,350 @@
+def convert_sd(sd, cpu=False):
+    sd_converted = {}
+    mapping = get_mapping()
+    for k, v in sd.items():
+        sd_converted[mapping[k]] = v.cpu() if cpu else v
+    return sd_converted
+
+
+def get_mapping():
+    return {
+        "encoder.project_in.conv.bias": "encoder.project_in.bias",
+        "encoder.project_in.conv.weight": "encoder.project_in.weight",
+        "encoder.stages.0.op_list.0.main.conv1.conv.bias": "encoder.stages.0.0.conv1.conv.bias",
+        "encoder.stages.0.op_list.0.main.conv1.conv.weight": "encoder.stages.0.0.conv1.conv.weight",
+        "encoder.stages.0.op_list.0.main.conv2.conv.weight": "encoder.stages.0.0.conv2.conv.weight",
+        "encoder.stages.0.op_list.0.main.conv2.norm.bias": "encoder.stages.0.0.conv2.norm.bias",
+        "encoder.stages.0.op_list.0.main.conv2.norm.weight": "encoder.stages.0.0.conv2.norm.weight",
+        "encoder.stages.0.op_list.1.main.conv1.conv.bias": "encoder.stages.0.1.conv1.conv.bias",
+        "encoder.stages.0.op_list.1.main.conv1.conv.weight": "encoder.stages.0.1.conv1.conv.weight",
+        "encoder.stages.0.op_list.1.main.conv2.conv.weight": "encoder.stages.0.1.conv2.conv.weight",
+        "encoder.stages.0.op_list.1.main.conv2.norm.bias": "encoder.stages.0.1.conv2.norm.bias",
+        "encoder.stages.0.op_list.1.main.conv2.norm.weight": "encoder.stages.0.1.conv2.norm.weight",
+        "encoder.stages.0.op_list.2.main.conv.bias": "encoder.stages.0.2.main.bias",
+        "encoder.stages.0.op_list.2.main.conv.weight": "encoder.stages.0.2.main.weight",
+        "encoder.stages.1.op_list.0.main.conv1.conv.bias": "encoder.stages.1.0.conv1.conv.bias",
+        "encoder.stages.1.op_list.0.main.conv1.conv.weight": "encoder.stages.1.0.conv1.conv.weight",
+        "encoder.stages.1.op_list.0.main.conv2.conv.weight": "encoder.stages.1.0.conv2.conv.weight",
+        "encoder.stages.1.op_list.0.main.conv2.norm.bias": "encoder.stages.1.0.conv2.norm.bias",
+        "encoder.stages.1.op_list.0.main.conv2.norm.weight": "encoder.stages.1.0.conv2.norm.weight",
+        "encoder.stages.1.op_list.1.main.conv1.conv.bias": "encoder.stages.1.1.conv1.conv.bias",
+        "encoder.stages.1.op_list.1.main.conv1.conv.weight": "encoder.stages.1.1.conv1.conv.weight",
+        "encoder.stages.1.op_list.1.main.conv2.conv.weight": "encoder.stages.1.1.conv2.conv.weight",
+        "encoder.stages.1.op_list.1.main.conv2.norm.bias": "encoder.stages.1.1.conv2.norm.bias",
+        "encoder.stages.1.op_list.1.main.conv2.norm.weight": "encoder.stages.1.1.conv2.norm.weight",
+        "encoder.stages.1.op_list.2.main.conv.bias": "encoder.stages.1.2.main.bias",
+        "encoder.stages.1.op_list.2.main.conv.weight": "encoder.stages.1.2.main.weight",
+        "encoder.stages.2.op_list.0.main.conv1.conv.bias": "encoder.stages.2.0.conv1.conv.bias",
+        "encoder.stages.2.op_list.0.main.conv1.conv.weight": "encoder.stages.2.0.conv1.conv.weight",
+        "encoder.stages.2.op_list.0.main.conv2.conv.weight": "encoder.stages.2.0.conv2.conv.weight",
+        "encoder.stages.2.op_list.0.main.conv2.norm.bias": "encoder.stages.2.0.conv2.norm.bias",
+        "encoder.stages.2.op_list.0.main.conv2.norm.weight": "encoder.stages.2.0.conv2.norm.weight",
+        "encoder.stages.2.op_list.1.main.conv1.conv.bias": "encoder.stages.2.1.conv1.conv.bias",
+        "encoder.stages.2.op_list.1.main.conv1.conv.weight": "encoder.stages.2.1.conv1.conv.weight",
+        "encoder.stages.2.op_list.1.main.conv2.conv.weight": "encoder.stages.2.1.conv2.conv.weight",
+        "encoder.stages.2.op_list.1.main.conv2.norm.bias": "encoder.stages.2.1.conv2.norm.bias",
+        "encoder.stages.2.op_list.1.main.conv2.norm.weight": "encoder.stages.2.1.conv2.norm.weight",
+        "encoder.stages.2.op_list.2.main.conv.bias": "encoder.stages.2.2.main.bias",
+        "encoder.stages.2.op_list.2.main.conv.weight": "encoder.stages.2.2.main.weight",
+        "encoder.stages.3.op_list.0.context_module.main.aggreg.0.0.weight": "encoder.stages.3.0.context_module.aggreg.0.0.weight",
+        "encoder.stages.3.op_list.0.context_module.main.aggreg.0.1.weight": "encoder.stages.3.0.context_module.aggreg.0.1.weight",
+        "encoder.stages.3.op_list.0.context_module.main.proj.conv.weight": "encoder.stages.3.0.context_module.proj.0.weight",
+        "encoder.stages.3.op_list.0.context_module.main.proj.norm.bias": "encoder.stages.3.0.context_module.proj.1.bias",
+        "encoder.stages.3.op_list.0.context_module.main.proj.norm.weight": "encoder.stages.3.0.context_module.proj.1.weight",
+        "encoder.stages.3.op_list.0.context_module.main.qkv.conv.weight": "encoder.stages.3.0.context_module.qkv.0.weight",
+        "encoder.stages.3.op_list.0.local_module.main.depth_conv.conv.bias": "encoder.stages.3.0.local_module.depth_conv.conv.bias",
+        "encoder.stages.3.op_list.0.local_module.main.depth_conv.conv.weight": "encoder.stages.3.0.local_module.depth_conv.conv.weight",
+        "encoder.stages.3.op_list.0.local_module.main.inverted_conv.conv.bias": "encoder.stages.3.0.local_module.inverted_conv.conv.bias",
+        "encoder.stages.3.op_list.0.local_module.main.inverted_conv.conv.weight": "encoder.stages.3.0.local_module.inverted_conv.conv.weight",
+        "encoder.stages.3.op_list.0.local_module.main.point_conv.conv.weight": "encoder.stages.3.0.local_module.point_conv.conv.weight",
+        "encoder.stages.3.op_list.0.local_module.main.point_conv.norm.bias": "encoder.stages.3.0.local_module.point_conv.norm.bias",
+        "encoder.stages.3.op_list.0.local_module.main.point_conv.norm.weight": "encoder.stages.3.0.local_module.point_conv.norm.weight",
+        "encoder.stages.3.op_list.1.context_module.main.aggreg.0.0.weight": "encoder.stages.3.1.context_module.aggreg.0.0.weight",
+        "encoder.stages.3.op_list.1.context_module.main.aggreg.0.1.weight": "encoder.stages.3.1.context_module.aggreg.0.1.weight",
+        "encoder.stages.3.op_list.1.context_module.main.proj.conv.weight": "encoder.stages.3.1.context_module.proj.0.weight",
+        "encoder.stages.3.op_list.1.context_module.main.proj.norm.bias": "encoder.stages.3.1.context_module.proj.1.bias",
+        "encoder.stages.3.op_list.1.context_module.main.proj.norm.weight": "encoder.stages.3.1.context_module.proj.1.weight",
+        "encoder.stages.3.op_list.1.context_module.main.qkv.conv.weight": "encoder.stages.3.1.context_module.qkv.0.weight",
+        "encoder.stages.3.op_list.1.local_module.main.depth_conv.conv.bias": "encoder.stages.3.1.local_module.depth_conv.conv.bias",
+        "encoder.stages.3.op_list.1.local_module.main.depth_conv.conv.weight": "encoder.stages.3.1.local_module.depth_conv.conv.weight",
+        "encoder.stages.3.op_list.1.local_module.main.inverted_conv.conv.bias": "encoder.stages.3.1.local_module.inverted_conv.conv.bias",
+        "encoder.stages.3.op_list.1.local_module.main.inverted_conv.conv.weight": "encoder.stages.3.1.local_module.inverted_conv.conv.weight",
+        "encoder.stages.3.op_list.1.local_module.main.point_conv.conv.weight": "encoder.stages.3.1.local_module.point_conv.conv.weight",
+        "encoder.stages.3.op_list.1.local_module.main.point_conv.norm.bias": "encoder.stages.3.1.local_module.point_conv.norm.bias",
+        "encoder.stages.3.op_list.1.local_module.main.point_conv.norm.weight": "encoder.stages.3.1.local_module.point_conv.norm.weight",
+        "encoder.stages.3.op_list.2.context_module.main.aggreg.0.0.weight": "encoder.stages.3.2.context_module.aggreg.0.0.weight",
+        "encoder.stages.3.op_list.2.context_module.main.aggreg.0.1.weight": "encoder.stages.3.2.context_module.aggreg.0.1.weight",
+        "encoder.stages.3.op_list.2.context_module.main.proj.conv.weight": "encoder.stages.3.2.context_module.proj.0.weight",
+        "encoder.stages.3.op_list.2.context_module.main.proj.norm.bias": "encoder.stages.3.2.context_module.proj.1.bias",
+        "encoder.stages.3.op_list.2.context_module.main.proj.norm.weight": "encoder.stages.3.2.context_module.proj.1.weight",
+        "encoder.stages.3.op_list.2.context_module.main.qkv.conv.weight": "encoder.stages.3.2.context_module.qkv.0.weight",
+        "encoder.stages.3.op_list.2.local_module.main.depth_conv.conv.bias": "encoder.stages.3.2.local_module.depth_conv.conv.bias",
+        "encoder.stages.3.op_list.2.local_module.main.depth_conv.conv.weight": "encoder.stages.3.2.local_module.depth_conv.conv.weight",
+        "encoder.stages.3.op_list.2.local_module.main.inverted_conv.conv.bias": "encoder.stages.3.2.local_module.inverted_conv.conv.bias",
+        "encoder.stages.3.op_list.2.local_module.main.inverted_conv.conv.weight": "encoder.stages.3.2.local_module.inverted_conv.conv.weight",
+        "encoder.stages.3.op_list.2.local_module.main.point_conv.conv.weight": "encoder.stages.3.2.local_module.point_conv.conv.weight",
+        "encoder.stages.3.op_list.2.local_module.main.point_conv.norm.bias": "encoder.stages.3.2.local_module.point_conv.norm.bias",
+        "encoder.stages.3.op_list.2.local_module.main.point_conv.norm.weight": "encoder.stages.3.2.local_module.point_conv.norm.weight",
+        "encoder.stages.3.op_list.3.main.conv.bias": "encoder.stages.3.3.main.bias",
+        "encoder.stages.3.op_list.3.main.conv.weight": "encoder.stages.3.3.main.weight",
+        "encoder.stages.4.op_list.0.context_module.main.aggreg.0.0.weight": "encoder.stages.4.0.context_module.aggreg.0.0.weight",
+        "encoder.stages.4.op_list.0.context_module.main.aggreg.0.1.weight": "encoder.stages.4.0.context_module.aggreg.0.1.weight",
+        "encoder.stages.4.op_list.0.context_module.main.proj.conv.weight": "encoder.stages.4.0.context_module.proj.0.weight",
+        "encoder.stages.4.op_list.0.context_module.main.proj.norm.bias": "encoder.stages.4.0.context_module.proj.1.bias",
+        "encoder.stages.4.op_list.0.context_module.main.proj.norm.weight": "encoder.stages.4.0.context_module.proj.1.weight",
+        "encoder.stages.4.op_list.0.context_module.main.qkv.conv.weight": "encoder.stages.4.0.context_module.qkv.0.weight",
+        "encoder.stages.4.op_list.0.local_module.main.depth_conv.conv.bias": "encoder.stages.4.0.local_module.depth_conv.conv.bias",
+        "encoder.stages.4.op_list.0.local_module.main.depth_conv.conv.weight": "encoder.stages.4.0.local_module.depth_conv.conv.weight",
+        "encoder.stages.4.op_list.0.local_module.main.inverted_conv.conv.bias": "encoder.stages.4.0.local_module.inverted_conv.conv.bias",
+        "encoder.stages.4.op_list.0.local_module.main.inverted_conv.conv.weight": "encoder.stages.4.0.local_module.inverted_conv.conv.weight",
+        "encoder.stages.4.op_list.0.local_module.main.point_conv.conv.weight": "encoder.stages.4.0.local_module.point_conv.conv.weight",
+        "encoder.stages.4.op_list.0.local_module.main.point_conv.norm.bias": "encoder.stages.4.0.local_module.point_conv.norm.bias",
+        "encoder.stages.4.op_list.0.local_module.main.point_conv.norm.weight": "encoder.stages.4.0.local_module.point_conv.norm.weight",
+        "encoder.stages.4.op_list.1.context_module.main.aggreg.0.0.weight": "encoder.stages.4.1.context_module.aggreg.0.0.weight",
+        "encoder.stages.4.op_list.1.context_module.main.aggreg.0.1.weight": "encoder.stages.4.1.context_module.aggreg.0.1.weight",
+        "encoder.stages.4.op_list.1.context_module.main.proj.conv.weight": "encoder.stages.4.1.context_module.proj.0.weight",
+        "encoder.stages.4.op_list.1.context_module.main.proj.norm.bias": "encoder.stages.4.1.context_module.proj.1.bias",
+        "encoder.stages.4.op_list.1.context_module.main.proj.norm.weight": "encoder.stages.4.1.context_module.proj.1.weight",
+        "encoder.stages.4.op_list.1.context_module.main.qkv.conv.weight": "encoder.stages.4.1.context_module.qkv.0.weight",
+        "encoder.stages.4.op_list.1.local_module.main.depth_conv.conv.bias": "encoder.stages.4.1.local_module.depth_conv.conv.bias",
+        "encoder.stages.4.op_list.1.local_module.main.depth_conv.conv.weight": "encoder.stages.4.1.local_module.depth_conv.conv.weight",
+        "encoder.stages.4.op_list.1.local_module.main.inverted_conv.conv.bias": "encoder.stages.4.1.local_module.inverted_conv.conv.bias",
+        "encoder.stages.4.op_list.1.local_module.main.inverted_conv.conv.weight": "encoder.stages.4.1.local_module.inverted_conv.conv.weight",
+        "encoder.stages.4.op_list.1.local_module.main.point_conv.conv.weight": "encoder.stages.4.1.local_module.point_conv.conv.weight",
+        "encoder.stages.4.op_list.1.local_module.main.point_conv.norm.bias": "encoder.stages.4.1.local_module.point_conv.norm.bias",
+        "encoder.stages.4.op_list.1.local_module.main.point_conv.norm.weight": "encoder.stages.4.1.local_module.point_conv.norm.weight",
+        "encoder.stages.4.op_list.2.context_module.main.aggreg.0.0.weight": "encoder.stages.4.2.context_module.aggreg.0.0.weight",
+        "encoder.stages.4.op_list.2.context_module.main.aggreg.0.1.weight": "encoder.stages.4.2.context_module.aggreg.0.1.weight",
+        "encoder.stages.4.op_list.2.context_module.main.proj.conv.weight": "encoder.stages.4.2.context_module.proj.0.weight",
+        "encoder.stages.4.op_list.2.context_module.main.proj.norm.bias": "encoder.stages.4.2.context_module.proj.1.bias",
+        "encoder.stages.4.op_list.2.context_module.main.proj.norm.weight": "encoder.stages.4.2.context_module.proj.1.weight",
+        "encoder.stages.4.op_list.2.context_module.main.qkv.conv.weight": "encoder.stages.4.2.context_module.qkv.0.weight",
+        "encoder.stages.4.op_list.2.local_module.main.depth_conv.conv.bias": "encoder.stages.4.2.local_module.depth_conv.conv.bias",
+        "encoder.stages.4.op_list.2.local_module.main.depth_conv.conv.weight": "encoder.stages.4.2.local_module.depth_conv.conv.weight",
+        "encoder.stages.4.op_list.2.local_module.main.inverted_conv.conv.bias": "encoder.stages.4.2.local_module.inverted_conv.conv.bias",
+        "encoder.stages.4.op_list.2.local_module.main.inverted_conv.conv.weight": "encoder.stages.4.2.local_module.inverted_conv.conv.weight",
+        "encoder.stages.4.op_list.2.local_module.main.point_conv.conv.weight": "encoder.stages.4.2.local_module.point_conv.conv.weight",
+        "encoder.stages.4.op_list.2.local_module.main.point_conv.norm.bias": "encoder.stages.4.2.local_module.point_conv.norm.bias",
+        "encoder.stages.4.op_list.2.local_module.main.point_conv.norm.weight": "encoder.stages.4.2.local_module.point_conv.norm.weight",
+        "encoder.stages.4.op_list.3.main.conv.bias": "encoder.stages.4.3.main.bias",
+        "encoder.stages.4.op_list.3.main.conv.weight": "encoder.stages.4.3.main.weight",
+        "encoder.stages.5.op_list.0.context_module.main.aggreg.0.0.weight": "encoder.stages.5.0.context_module.aggreg.0.0.weight",
+        "encoder.stages.5.op_list.0.context_module.main.aggreg.0.1.weight": "encoder.stages.5.0.context_module.aggreg.0.1.weight",
+        "encoder.stages.5.op_list.0.context_module.main.proj.conv.weight": "encoder.stages.5.0.context_module.proj.0.weight",
+        "encoder.stages.5.op_list.0.context_module.main.proj.norm.bias": "encoder.stages.5.0.context_module.proj.1.bias",
+        "encoder.stages.5.op_list.0.context_module.main.proj.norm.weight": "encoder.stages.5.0.context_module.proj.1.weight",
+        "encoder.stages.5.op_list.0.context_module.main.qkv.conv.weight": "encoder.stages.5.0.context_module.qkv.0.weight",
+        "encoder.stages.5.op_list.0.local_module.main.depth_conv.conv.bias": "encoder.stages.5.0.local_module.depth_conv.conv.bias",
+        "encoder.stages.5.op_list.0.local_module.main.depth_conv.conv.weight": "encoder.stages.5.0.local_module.depth_conv.conv.weight",
+        "encoder.stages.5.op_list.0.local_module.main.inverted_conv.conv.bias": "encoder.stages.5.0.local_module.inverted_conv.conv.bias",
+        "encoder.stages.5.op_list.0.local_module.main.inverted_conv.conv.weight": "encoder.stages.5.0.local_module.inverted_conv.conv.weight",
+        "encoder.stages.5.op_list.0.local_module.main.point_conv.conv.weight": "encoder.stages.5.0.local_module.point_conv.conv.weight",
+        "encoder.stages.5.op_list.0.local_module.main.point_conv.norm.bias": "encoder.stages.5.0.local_module.point_conv.norm.bias",
+        "encoder.stages.5.op_list.0.local_module.main.point_conv.norm.weight": "encoder.stages.5.0.local_module.point_conv.norm.weight",
+        "encoder.stages.5.op_list.1.context_module.main.aggreg.0.0.weight": "encoder.stages.5.1.context_module.aggreg.0.0.weight",
+        "encoder.stages.5.op_list.1.context_module.main.aggreg.0.1.weight": "encoder.stages.5.1.context_module.aggreg.0.1.weight",
+        "encoder.stages.5.op_list.1.context_module.main.proj.conv.weight": "encoder.stages.5.1.context_module.proj.0.weight",
+        "encoder.stages.5.op_list.1.context_module.main.proj.norm.bias": "encoder.stages.5.1.context_module.proj.1.bias",
+        "encoder.stages.5.op_list.1.context_module.main.proj.norm.weight": "encoder.stages.5.1.context_module.proj.1.weight",
+        "encoder.stages.5.op_list.1.context_module.main.qkv.conv.weight": "encoder.stages.5.1.context_module.qkv.0.weight",
+        "encoder.stages.5.op_list.1.local_module.main.depth_conv.conv.bias": "encoder.stages.5.1.local_module.depth_conv.conv.bias",
+        "encoder.stages.5.op_list.1.local_module.main.depth_conv.conv.weight": "encoder.stages.5.1.local_module.depth_conv.conv.weight",
+        "encoder.stages.5.op_list.1.local_module.main.inverted_conv.conv.bias": "encoder.stages.5.1.local_module.inverted_conv.conv.bias",
+        "encoder.stages.5.op_list.1.local_module.main.inverted_conv.conv.weight": "encoder.stages.5.1.local_module.inverted_conv.conv.weight",
+        "encoder.stages.5.op_list.1.local_module.main.point_conv.conv.weight": "encoder.stages.5.1.local_module.point_conv.conv.weight",
+        "encoder.stages.5.op_list.1.local_module.main.point_conv.norm.bias": "encoder.stages.5.1.local_module.point_conv.norm.bias",
+        "encoder.stages.5.op_list.1.local_module.main.point_conv.norm.weight": "encoder.stages.5.1.local_module.point_conv.norm.weight",
+        "encoder.stages.5.op_list.2.context_module.main.aggreg.0.0.weight": "encoder.stages.5.2.context_module.aggreg.0.0.weight",
+        "encoder.stages.5.op_list.2.context_module.main.aggreg.0.1.weight": "encoder.stages.5.2.context_module.aggreg.0.1.weight",
+        "encoder.stages.5.op_list.2.context_module.main.proj.conv.weight": "encoder.stages.5.2.context_module.proj.0.weight",
+        "encoder.stages.5.op_list.2.context_module.main.proj.norm.bias": "encoder.stages.5.2.context_module.proj.1.bias",
+        "encoder.stages.5.op_list.2.context_module.main.proj.norm.weight": "encoder.stages.5.2.context_module.proj.1.weight",
+        "encoder.stages.5.op_list.2.context_module.main.qkv.conv.weight": "encoder.stages.5.2.context_module.qkv.0.weight",
+        "encoder.stages.5.op_list.2.local_module.main.depth_conv.conv.bias": "encoder.stages.5.2.local_module.depth_conv.conv.bias",
+        "encoder.stages.5.op_list.2.local_module.main.depth_conv.conv.weight": "encoder.stages.5.2.local_module.depth_conv.conv.weight",
+        "encoder.stages.5.op_list.2.local_module.main.inverted_conv.conv.bias": "encoder.stages.5.2.local_module.inverted_conv.conv.bias",
+        "encoder.stages.5.op_list.2.local_module.main.inverted_conv.conv.weight": "encoder.stages.5.2.local_module.inverted_conv.conv.weight",
+        "encoder.stages.5.op_list.2.local_module.main.point_conv.conv.weight": "encoder.stages.5.2.local_module.point_conv.conv.weight",
+        "encoder.stages.5.op_list.2.local_module.main.point_conv.norm.bias": "encoder.stages.5.2.local_module.point_conv.norm.bias",
+        "encoder.stages.5.op_list.2.local_module.main.point_conv.norm.weight": "encoder.stages.5.2.local_module.point_conv.norm.weight",
+        "encoder.project_out.main.op_list.0.conv.bias": "encoder.project_out.main.0.conv.bias",
+        "encoder.project_out.main.op_list.0.conv.weight": "encoder.project_out.main.0.conv.weight",
+        "decoder.project_in.main.conv.bias": "decoder.project_in.main.conv.bias",
+        "decoder.project_in.main.conv.weight": "decoder.project_in.main.conv.weight",
+        "decoder.stages.0.op_list.0.main.conv.conv.bias": "decoder.stages.0.0.main.conv.bias",
+        "decoder.stages.0.op_list.0.main.conv.conv.weight": "decoder.stages.0.0.main.conv.weight",
+        "decoder.stages.0.op_list.1.main.conv1.conv.bias": "decoder.stages.0.1.conv1.conv.bias",
+        "decoder.stages.0.op_list.1.main.conv1.conv.weight": "decoder.stages.0.1.conv1.conv.weight",
+        "decoder.stages.0.op_list.1.main.conv2.conv.weight": "decoder.stages.0.1.conv2.conv.weight",
+        "decoder.stages.0.op_list.1.main.conv2.norm.bias": "decoder.stages.0.1.conv2.norm.bias",
+        "decoder.stages.0.op_list.1.main.conv2.norm.weight": "decoder.stages.0.1.conv2.norm.weight",
+        "decoder.stages.0.op_list.2.main.conv1.conv.bias": "decoder.stages.0.2.conv1.conv.bias",
+        "decoder.stages.0.op_list.2.main.conv1.conv.weight": "decoder.stages.0.2.conv1.conv.weight",
+        "decoder.stages.0.op_list.2.main.conv2.conv.weight": "decoder.stages.0.2.conv2.conv.weight",
+        "decoder.stages.0.op_list.2.main.conv2.norm.bias": "decoder.stages.0.2.conv2.norm.bias",
+        "decoder.stages.0.op_list.2.main.conv2.norm.weight": "decoder.stages.0.2.conv2.norm.weight",
+        "decoder.stages.0.op_list.3.main.conv1.conv.bias": "decoder.stages.0.3.conv1.conv.bias",
+        "decoder.stages.0.op_list.3.main.conv1.conv.weight": "decoder.stages.0.3.conv1.conv.weight",
+        "decoder.stages.0.op_list.3.main.conv2.conv.weight": "decoder.stages.0.3.conv2.conv.weight",
+        "decoder.stages.0.op_list.3.main.conv2.norm.bias": "decoder.stages.0.3.conv2.norm.bias",
+        "decoder.stages.0.op_list.3.main.conv2.norm.weight": "decoder.stages.0.3.conv2.norm.weight",
+        "decoder.stages.1.op_list.0.main.conv.conv.bias": "decoder.stages.1.0.main.conv.bias",
+        "decoder.stages.1.op_list.0.main.conv.conv.weight": "decoder.stages.1.0.main.conv.weight",
+        "decoder.stages.1.op_list.1.main.conv1.conv.bias": "decoder.stages.1.1.conv1.conv.bias",
+        "decoder.stages.1.op_list.1.main.conv1.conv.weight": "decoder.stages.1.1.conv1.conv.weight",
+        "decoder.stages.1.op_list.1.main.conv2.conv.weight": "decoder.stages.1.1.conv2.conv.weight",
+        "decoder.stages.1.op_list.1.main.conv2.norm.bias": "decoder.stages.1.1.conv2.norm.bias",
+        "decoder.stages.1.op_list.1.main.conv2.norm.weight": "decoder.stages.1.1.conv2.norm.weight",
+        "decoder.stages.1.op_list.2.main.conv1.conv.bias": "decoder.stages.1.2.conv1.conv.bias",
+        "decoder.stages.1.op_list.2.main.conv1.conv.weight": "decoder.stages.1.2.conv1.conv.weight",
+        "decoder.stages.1.op_list.2.main.conv2.conv.weight": "decoder.stages.1.2.conv2.conv.weight",
+        "decoder.stages.1.op_list.2.main.conv2.norm.bias": "decoder.stages.1.2.conv2.norm.bias",
+        "decoder.stages.1.op_list.2.main.conv2.norm.weight": "decoder.stages.1.2.conv2.norm.weight",
+        "decoder.stages.1.op_list.3.main.conv1.conv.bias": "decoder.stages.1.3.conv1.conv.bias",
+        "decoder.stages.1.op_list.3.main.conv1.conv.weight": "decoder.stages.1.3.conv1.conv.weight",
+        "decoder.stages.1.op_list.3.main.conv2.conv.weight": "decoder.stages.1.3.conv2.conv.weight",
+        "decoder.stages.1.op_list.3.main.conv2.norm.bias": "decoder.stages.1.3.conv2.norm.bias",
+        "decoder.stages.1.op_list.3.main.conv2.norm.weight": "decoder.stages.1.3.conv2.norm.weight",
+        "decoder.stages.2.op_list.0.main.conv.conv.bias": "decoder.stages.2.0.main.conv.bias",
+        "decoder.stages.2.op_list.0.main.conv.conv.weight": "decoder.stages.2.0.main.conv.weight",
+        "decoder.stages.2.op_list.1.main.conv1.conv.bias": "decoder.stages.2.1.conv1.conv.bias",
+        "decoder.stages.2.op_list.1.main.conv1.conv.weight": "decoder.stages.2.1.conv1.conv.weight",
+        "decoder.stages.2.op_list.1.main.conv2.conv.weight": "decoder.stages.2.1.conv2.conv.weight",
+        "decoder.stages.2.op_list.1.main.conv2.norm.bias": "decoder.stages.2.1.conv2.norm.bias",
+        "decoder.stages.2.op_list.1.main.conv2.norm.weight": "decoder.stages.2.1.conv2.norm.weight",
+        "decoder.stages.2.op_list.2.main.conv1.conv.bias": "decoder.stages.2.2.conv1.conv.bias",
+        "decoder.stages.2.op_list.2.main.conv1.conv.weight": "decoder.stages.2.2.conv1.conv.weight",
+        "decoder.stages.2.op_list.2.main.conv2.conv.weight": "decoder.stages.2.2.conv2.conv.weight",
+        "decoder.stages.2.op_list.2.main.conv2.norm.bias": "decoder.stages.2.2.conv2.norm.bias",
+        "decoder.stages.2.op_list.2.main.conv2.norm.weight": "decoder.stages.2.2.conv2.norm.weight",
+        "decoder.stages.2.op_list.3.main.conv1.conv.bias": "decoder.stages.2.3.conv1.conv.bias",
+        "decoder.stages.2.op_list.3.main.conv1.conv.weight": "decoder.stages.2.3.conv1.conv.weight",
+        "decoder.stages.2.op_list.3.main.conv2.conv.weight": "decoder.stages.2.3.conv2.conv.weight",
+        "decoder.stages.2.op_list.3.main.conv2.norm.bias": "decoder.stages.2.3.conv2.norm.bias",
+        "decoder.stages.2.op_list.3.main.conv2.norm.weight": "decoder.stages.2.3.conv2.norm.weight",
+        "decoder.stages.3.op_list.0.main.conv.conv.bias": "decoder.stages.3.0.main.conv.bias",
+        "decoder.stages.3.op_list.0.main.conv.conv.weight": "decoder.stages.3.0.main.conv.weight",
+        "decoder.stages.3.op_list.1.context_module.main.aggreg.0.0.weight": "decoder.stages.3.1.context_module.aggreg.0.0.weight",
+        "decoder.stages.3.op_list.1.context_module.main.aggreg.0.1.weight": "decoder.stages.3.1.context_module.aggreg.0.1.weight",
+        "decoder.stages.3.op_list.1.context_module.main.proj.conv.weight": "decoder.stages.3.1.context_module.proj.0.weight",
+        "decoder.stages.3.op_list.1.context_module.main.proj.norm.bias": "decoder.stages.3.1.context_module.proj.1.bias",
+        "decoder.stages.3.op_list.1.context_module.main.proj.norm.weight": "decoder.stages.3.1.context_module.proj.1.weight",
+        "decoder.stages.3.op_list.1.context_module.main.qkv.conv.weight": "decoder.stages.3.1.context_module.qkv.0.weight",
+        "decoder.stages.3.op_list.1.local_module.main.depth_conv.conv.bias": "decoder.stages.3.1.local_module.depth_conv.conv.bias",
+        "decoder.stages.3.op_list.1.local_module.main.depth_conv.conv.weight": "decoder.stages.3.1.local_module.depth_conv.conv.weight",
+        "decoder.stages.3.op_list.1.local_module.main.inverted_conv.conv.bias": "decoder.stages.3.1.local_module.inverted_conv.conv.bias",
+        "decoder.stages.3.op_list.1.local_module.main.inverted_conv.conv.weight": "decoder.stages.3.1.local_module.inverted_conv.conv.weight",
+        "decoder.stages.3.op_list.1.local_module.main.point_conv.conv.weight": "decoder.stages.3.1.local_module.point_conv.conv.weight",
+        "decoder.stages.3.op_list.1.local_module.main.point_conv.norm.bias": "decoder.stages.3.1.local_module.point_conv.norm.bias",
+        "decoder.stages.3.op_list.1.local_module.main.point_conv.norm.weight": "decoder.stages.3.1.local_module.point_conv.norm.weight",
+        "decoder.stages.3.op_list.2.context_module.main.aggreg.0.0.weight": "decoder.stages.3.2.context_module.aggreg.0.0.weight",
+        "decoder.stages.3.op_list.2.context_module.main.aggreg.0.1.weight": "decoder.stages.3.2.context_module.aggreg.0.1.weight",
+        "decoder.stages.3.op_list.2.context_module.main.proj.conv.weight": "decoder.stages.3.2.context_module.proj.0.weight",
+        "decoder.stages.3.op_list.2.context_module.main.proj.norm.bias": "decoder.stages.3.2.context_module.proj.1.bias",
+        "decoder.stages.3.op_list.2.context_module.main.proj.norm.weight": "decoder.stages.3.2.context_module.proj.1.weight",
+        "decoder.stages.3.op_list.2.context_module.main.qkv.conv.weight": "decoder.stages.3.2.context_module.qkv.0.weight",
+        "decoder.stages.3.op_list.2.local_module.main.depth_conv.conv.bias": "decoder.stages.3.2.local_module.depth_conv.conv.bias",
+        "decoder.stages.3.op_list.2.local_module.main.depth_conv.conv.weight": "decoder.stages.3.2.local_module.depth_conv.conv.weight",
+        "decoder.stages.3.op_list.2.local_module.main.inverted_conv.conv.bias": "decoder.stages.3.2.local_module.inverted_conv.conv.bias",
+        "decoder.stages.3.op_list.2.local_module.main.inverted_conv.conv.weight": "decoder.stages.3.2.local_module.inverted_conv.conv.weight",
+        "decoder.stages.3.op_list.2.local_module.main.point_conv.conv.weight": "decoder.stages.3.2.local_module.point_conv.conv.weight",
+        "decoder.stages.3.op_list.2.local_module.main.point_conv.norm.bias": "decoder.stages.3.2.local_module.point_conv.norm.bias",
+        "decoder.stages.3.op_list.2.local_module.main.point_conv.norm.weight": "decoder.stages.3.2.local_module.point_conv.norm.weight",
+        "decoder.stages.3.op_list.3.context_module.main.aggreg.0.0.weight": "decoder.stages.3.3.context_module.aggreg.0.0.weight",
+        "decoder.stages.3.op_list.3.context_module.main.aggreg.0.1.weight": "decoder.stages.3.3.context_module.aggreg.0.1.weight",
+        "decoder.stages.3.op_list.3.context_module.main.proj.conv.weight": "decoder.stages.3.3.context_module.proj.0.weight",
+        "decoder.stages.3.op_list.3.context_module.main.proj.norm.bias": "decoder.stages.3.3.context_module.proj.1.bias",
+        "decoder.stages.3.op_list.3.context_module.main.proj.norm.weight": "decoder.stages.3.3.context_module.proj.1.weight",
+        "decoder.stages.3.op_list.3.context_module.main.qkv.conv.weight": "decoder.stages.3.3.context_module.qkv.0.weight",
+        "decoder.stages.3.op_list.3.local_module.main.depth_conv.conv.bias": "decoder.stages.3.3.local_module.depth_conv.conv.bias",
+        "decoder.stages.3.op_list.3.local_module.main.depth_conv.conv.weight": "decoder.stages.3.3.local_module.depth_conv.conv.weight",
+        "decoder.stages.3.op_list.3.local_module.main.inverted_conv.conv.bias": "decoder.stages.3.3.local_module.inverted_conv.conv.bias",
+        "decoder.stages.3.op_list.3.local_module.main.inverted_conv.conv.weight": "decoder.stages.3.3.local_module.inverted_conv.conv.weight",
+        "decoder.stages.3.op_list.3.local_module.main.point_conv.conv.weight": "decoder.stages.3.3.local_module.point_conv.conv.weight",
+        "decoder.stages.3.op_list.3.local_module.main.point_conv.norm.bias": "decoder.stages.3.3.local_module.point_conv.norm.bias",
+        "decoder.stages.3.op_list.3.local_module.main.point_conv.norm.weight": "decoder.stages.3.3.local_module.point_conv.norm.weight",
+        "decoder.stages.4.op_list.0.main.conv.conv.bias": "decoder.stages.4.0.main.conv.bias",
+        "decoder.stages.4.op_list.0.main.conv.conv.weight": "decoder.stages.4.0.main.conv.weight",
+        "decoder.stages.4.op_list.1.context_module.main.aggreg.0.0.weight": "decoder.stages.4.1.context_module.aggreg.0.0.weight",
+        "decoder.stages.4.op_list.1.context_module.main.aggreg.0.1.weight": "decoder.stages.4.1.context_module.aggreg.0.1.weight",
+        "decoder.stages.4.op_list.1.context_module.main.proj.conv.weight": "decoder.stages.4.1.context_module.proj.0.weight",
+        "decoder.stages.4.op_list.1.context_module.main.proj.norm.bias": "decoder.stages.4.1.context_module.proj.1.bias",
+        "decoder.stages.4.op_list.1.context_module.main.proj.norm.weight": "decoder.stages.4.1.context_module.proj.1.weight",
+        "decoder.stages.4.op_list.1.context_module.main.qkv.conv.weight": "decoder.stages.4.1.context_module.qkv.0.weight",
+        "decoder.stages.4.op_list.1.local_module.main.depth_conv.conv.bias": "decoder.stages.4.1.local_module.depth_conv.conv.bias",
+        "decoder.stages.4.op_list.1.local_module.main.depth_conv.conv.weight": "decoder.stages.4.1.local_module.depth_conv.conv.weight",
+        "decoder.stages.4.op_list.1.local_module.main.inverted_conv.conv.bias": "decoder.stages.4.1.local_module.inverted_conv.conv.bias",
+        "decoder.stages.4.op_list.1.local_module.main.inverted_conv.conv.weight": "decoder.stages.4.1.local_module.inverted_conv.conv.weight",
+        "decoder.stages.4.op_list.1.local_module.main.point_conv.conv.weight": "decoder.stages.4.1.local_module.point_conv.conv.weight",
+        "decoder.stages.4.op_list.1.local_module.main.point_conv.norm.bias": "decoder.stages.4.1.local_module.point_conv.norm.bias",
+        "decoder.stages.4.op_list.1.local_module.main.point_conv.norm.weight": "decoder.stages.4.1.local_module.point_conv.norm.weight",
+        "decoder.stages.4.op_list.2.context_module.main.aggreg.0.0.weight": "decoder.stages.4.2.context_module.aggreg.0.0.weight",
+        "decoder.stages.4.op_list.2.context_module.main.aggreg.0.1.weight": "decoder.stages.4.2.context_module.aggreg.0.1.weight",
+        "decoder.stages.4.op_list.2.context_module.main.proj.conv.weight": "decoder.stages.4.2.context_module.proj.0.weight",
+        "decoder.stages.4.op_list.2.context_module.main.proj.norm.bias": "decoder.stages.4.2.context_module.proj.1.bias",
+        "decoder.stages.4.op_list.2.context_module.main.proj.norm.weight": "decoder.stages.4.2.context_module.proj.1.weight",
+        "decoder.stages.4.op_list.2.context_module.main.qkv.conv.weight": "decoder.stages.4.2.context_module.qkv.0.weight",
+        "decoder.stages.4.op_list.2.local_module.main.depth_conv.conv.bias": "decoder.stages.4.2.local_module.depth_conv.conv.bias",
+        "decoder.stages.4.op_list.2.local_module.main.depth_conv.conv.weight": "decoder.stages.4.2.local_module.depth_conv.conv.weight",
+        "decoder.stages.4.op_list.2.local_module.main.inverted_conv.conv.bias": "decoder.stages.4.2.local_module.inverted_conv.conv.bias",
+        "decoder.stages.4.op_list.2.local_module.main.inverted_conv.conv.weight": "decoder.stages.4.2.local_module.inverted_conv.conv.weight",
+        "decoder.stages.4.op_list.2.local_module.main.point_conv.conv.weight": "decoder.stages.4.2.local_module.point_conv.conv.weight",
+        "decoder.stages.4.op_list.2.local_module.main.point_conv.norm.bias": "decoder.stages.4.2.local_module.point_conv.norm.bias",
+        "decoder.stages.4.op_list.2.local_module.main.point_conv.norm.weight": "decoder.stages.4.2.local_module.point_conv.norm.weight",
+        "decoder.stages.4.op_list.3.context_module.main.aggreg.0.0.weight": "decoder.stages.4.3.context_module.aggreg.0.0.weight",
+        "decoder.stages.4.op_list.3.context_module.main.aggreg.0.1.weight": "decoder.stages.4.3.context_module.aggreg.0.1.weight",
+        "decoder.stages.4.op_list.3.context_module.main.proj.conv.weight": "decoder.stages.4.3.context_module.proj.0.weight",
+        "decoder.stages.4.op_list.3.context_module.main.proj.norm.bias": "decoder.stages.4.3.context_module.proj.1.bias",
+        "decoder.stages.4.op_list.3.context_module.main.proj.norm.weight": "decoder.stages.4.3.context_module.proj.1.weight",
+        "decoder.stages.4.op_list.3.context_module.main.qkv.conv.weight": "decoder.stages.4.3.context_module.qkv.0.weight",
+        "decoder.stages.4.op_list.3.local_module.main.depth_conv.conv.bias": "decoder.stages.4.3.local_module.depth_conv.conv.bias",
+        "decoder.stages.4.op_list.3.local_module.main.depth_conv.conv.weight": "decoder.stages.4.3.local_module.depth_conv.conv.weight",
+        "decoder.stages.4.op_list.3.local_module.main.inverted_conv.conv.bias": "decoder.stages.4.3.local_module.inverted_conv.conv.bias",
+        "decoder.stages.4.op_list.3.local_module.main.inverted_conv.conv.weight": "decoder.stages.4.3.local_module.inverted_conv.conv.weight",
+        "decoder.stages.4.op_list.3.local_module.main.point_conv.conv.weight": "decoder.stages.4.3.local_module.point_conv.conv.weight",
+        "decoder.stages.4.op_list.3.local_module.main.point_conv.norm.bias": "decoder.stages.4.3.local_module.point_conv.norm.bias",
+        "decoder.stages.4.op_list.3.local_module.main.point_conv.norm.weight": "decoder.stages.4.3.local_module.point_conv.norm.weight",
+        "decoder.stages.5.op_list.0.context_module.main.aggreg.0.0.weight": "decoder.stages.5.0.context_module.aggreg.0.0.weight",
+        "decoder.stages.5.op_list.0.context_module.main.aggreg.0.1.weight": "decoder.stages.5.0.context_module.aggreg.0.1.weight",
+        "decoder.stages.5.op_list.0.context_module.main.proj.conv.weight": "decoder.stages.5.0.context_module.proj.0.weight",
+        "decoder.stages.5.op_list.0.context_module.main.proj.norm.bias": "decoder.stages.5.0.context_module.proj.1.bias",
+        "decoder.stages.5.op_list.0.context_module.main.proj.norm.weight": "decoder.stages.5.0.context_module.proj.1.weight",
+        "decoder.stages.5.op_list.0.context_module.main.qkv.conv.weight": "decoder.stages.5.0.context_module.qkv.0.weight",
+        "decoder.stages.5.op_list.0.local_module.main.depth_conv.conv.bias": "decoder.stages.5.0.local_module.depth_conv.conv.bias",
+        "decoder.stages.5.op_list.0.local_module.main.depth_conv.conv.weight": "decoder.stages.5.0.local_module.depth_conv.conv.weight",
+        "decoder.stages.5.op_list.0.local_module.main.inverted_conv.conv.bias": "decoder.stages.5.0.local_module.inverted_conv.conv.bias",
+        "decoder.stages.5.op_list.0.local_module.main.inverted_conv.conv.weight": "decoder.stages.5.0.local_module.inverted_conv.conv.weight",
+        "decoder.stages.5.op_list.0.local_module.main.point_conv.conv.weight": "decoder.stages.5.0.local_module.point_conv.conv.weight",
+        "decoder.stages.5.op_list.0.local_module.main.point_conv.norm.bias": "decoder.stages.5.0.local_module.point_conv.norm.bias",
+        "decoder.stages.5.op_list.0.local_module.main.point_conv.norm.weight": "decoder.stages.5.0.local_module.point_conv.norm.weight",
+        "decoder.stages.5.op_list.1.context_module.main.aggreg.0.0.weight": "decoder.stages.5.1.context_module.aggreg.0.0.weight",
+        "decoder.stages.5.op_list.1.context_module.main.aggreg.0.1.weight": "decoder.stages.5.1.context_module.aggreg.0.1.weight",
+        "decoder.stages.5.op_list.1.context_module.main.proj.conv.weight": "decoder.stages.5.1.context_module.proj.0.weight",
+        "decoder.stages.5.op_list.1.context_module.main.proj.norm.bias": "decoder.stages.5.1.context_module.proj.1.bias",
+        "decoder.stages.5.op_list.1.context_module.main.proj.norm.weight": "decoder.stages.5.1.context_module.proj.1.weight",
+        "decoder.stages.5.op_list.1.context_module.main.qkv.conv.weight": "decoder.stages.5.1.context_module.qkv.0.weight",
+        "decoder.stages.5.op_list.1.local_module.main.depth_conv.conv.bias": "decoder.stages.5.1.local_module.depth_conv.conv.bias",
+        "decoder.stages.5.op_list.1.local_module.main.depth_conv.conv.weight": "decoder.stages.5.1.local_module.depth_conv.conv.weight",
+        "decoder.stages.5.op_list.1.local_module.main.inverted_conv.conv.bias": "decoder.stages.5.1.local_module.inverted_conv.conv.bias",
+        "decoder.stages.5.op_list.1.local_module.main.inverted_conv.conv.weight": "decoder.stages.5.1.local_module.inverted_conv.conv.weight",
+        "decoder.stages.5.op_list.1.local_module.main.point_conv.conv.weight": "decoder.stages.5.1.local_module.point_conv.conv.weight",
+        "decoder.stages.5.op_list.1.local_module.main.point_conv.norm.bias": "decoder.stages.5.1.local_module.point_conv.norm.bias",
+        "decoder.stages.5.op_list.1.local_module.main.point_conv.norm.weight": "decoder.stages.5.1.local_module.point_conv.norm.weight",
+        "decoder.stages.5.op_list.2.context_module.main.aggreg.0.0.weight": "decoder.stages.5.2.context_module.aggreg.0.0.weight",
+        "decoder.stages.5.op_list.2.context_module.main.aggreg.0.1.weight": "decoder.stages.5.2.context_module.aggreg.0.1.weight",
+        "decoder.stages.5.op_list.2.context_module.main.proj.conv.weight": "decoder.stages.5.2.context_module.proj.0.weight",
+        "decoder.stages.5.op_list.2.context_module.main.proj.norm.bias": "decoder.stages.5.2.context_module.proj.1.bias",
+        "decoder.stages.5.op_list.2.context_module.main.proj.norm.weight": "decoder.stages.5.2.context_module.proj.1.weight",
+        "decoder.stages.5.op_list.2.context_module.main.qkv.conv.weight": "decoder.stages.5.2.context_module.qkv.0.weight",
+        "decoder.stages.5.op_list.2.local_module.main.depth_conv.conv.bias": "decoder.stages.5.2.local_module.depth_conv.conv.bias",
+        "decoder.stages.5.op_list.2.local_module.main.depth_conv.conv.weight": "decoder.stages.5.2.local_module.depth_conv.conv.weight",
+        "decoder.stages.5.op_list.2.local_module.main.inverted_conv.conv.bias": "decoder.stages.5.2.local_module.inverted_conv.conv.bias",
+        "decoder.stages.5.op_list.2.local_module.main.inverted_conv.conv.weight": "decoder.stages.5.2.local_module.inverted_conv.conv.weight",
+        "decoder.stages.5.op_list.2.local_module.main.point_conv.conv.weight": "decoder.stages.5.2.local_module.point_conv.conv.weight",
+        "decoder.stages.5.op_list.2.local_module.main.point_conv.norm.bias": "decoder.stages.5.2.local_module.point_conv.norm.bias",
+        "decoder.stages.5.op_list.2.local_module.main.point_conv.norm.weight": "decoder.stages.5.2.local_module.point_conv.norm.weight",
+        "decoder.project_out.op_list.0.bias": "decoder.project_out.0.bias",
+        "decoder.project_out.op_list.0.weight": "decoder.project_out.0.weight",
+        "decoder.project_out.op_list.2.conv.bias": "decoder.project_out.2.conv.bias",
+        "decoder.project_out.op_list.2.conv.weight": "decoder.project_out.2.conv.weight",
+    }

requirements.txt

@@ -4,3 +4,4 @@ transformers>=4.34.1
 accelerate>=0.23.0
 einops>=0.6.0
 protobuf>=3.20.3
+bitsandbytes>=0.41.0