code: def _convert_first_stage(
      d: tensor_dict_type, md: tensor_dict_type
  ) -> tensor_dict_type:
      d, md = map(shallow_copy_dict, [d, md])
      keys = list(d.keys())
      m_keys = list(md.keys())
      for k in [
          "first_stage.codebook.embedding.weight",
          "first_stage.to_embedding.weight",
          "first_stage.to_embedding.bias",
          "first_stage.from_embedding.weight",
          "first_stage.from_embedding.bias",
      ]:
          if k in m_keys:
              m_keys.remove(k)
              m_keys.append(k)
      start_idx = None
      end_idx = None
      for i, k in enumerate(keys):
          if start_idx is None and k.startswith(
              "first_stage_model.decoder.up"
          ):
              start_idx = i
          elif k == "first_stage_model.decoder.norm_out.weight":
              end_idx = i
      if end_idx is None:
          raise ValueError(
              "`first_stage_model.decoder.norm_out.weight` is missing"
          )
      before = keys[:start_idx]
      up = keys[start_idx:end_idx]
      after = keys[end_idx:]
      num_up = int(keys[end_idx - 1].split(".")[3]) + 1
      new_up = []
      mapping = shallow_copy_dict(d)
      local_key_mapping = {k: k for k in d}
      for k in up:
          mapping.pop(k)
      for i in reversed(range(num_up)):
          for k in up:
              if k.startswith(f"first_stage_model.decoder.up.{i}"):
                  ks = k.split(".")
                  ks[3] = str(num_up - int(ks[3]))
                  new_k = ".".join(ks)
                  new_up.append(new_k)
                  mapping[new_k] = d[k]
                  local_key_mapping[new_k] = k
      keys = before + new_up + after
  
      def _handle_attn(prefix: str) -> None:
          attn_keys: Dict[int, Dict[int, List[str]]] = {}
          all_attn_keys = []
          for key in keys:
              if key.startswith(prefix):
                  ks = key.split(".")
                  if ks[4] == "attn":
                      level = int(ks[3])
                      idx = int(ks[5])
                      attn_keys.setdefault(level, {}).setdefault(
                          idx, []
                      ).append(key)
                      all_attn_keys.append(key)
          for level, l_attn_keys in attn_keys.items():
              for idx, idx_attn_keys in l_attn_keys.items():
                  for i, key in enumerate(keys):
                      if key.startswith(
                          f"{prefix}.{level}.block.{idx + 1}"
                      ):
                          for attn_key in idx_attn_keys:
                              keys.remove(attn_key)
                              keys.insert(i, attn_key)
                              i += 1
                          break
  
      _handle_attn("first_stage_model.encoder.down")
      _handle_attn("first_stage_model.decoder.up")
      new_d = {}
      for k, mk in zip(keys, m_keys):
          v = mapping[k]
          mv = md[mk]
          assert (
              v.shape == mv.shape
          ), f"{k} ({v.shape}) != {mk} ({mv.shape})"
          new_d[mk] = v
          key_mapping_inv[mk] = local_key_mapping[k]
      return new_d
  
imports: import json
  import torch
  from torch import Tensor
  from typing import Dict
  from typing import List
  from typing import Union
  from typing import Optional
  from cftool.misc import shallow_copy_dict
  from cftool.array import tensor_dict_type
  from ..misc.toolkit import get_tensors
  from ..misc.toolkit import download_static
  from ..api.cv.diffusion import DiffusionAPI
label_desc: Use image node_id 2 for calling a global function with example usage: _convert_first_stage(d, md) and returns: new_d