From 15cd96d709ea2791ce97f0d9b7c744968aff85da Mon Sep 17 00:00:00 2001 From: Anubhav Singh Date: Sun, 30 Aug 2026 22:12:45 +0530 Subject: [PATCH] Add CogVideoX Diffusers checkpoint conversion --- scripts/convert_diffusers_to_cogvideox.py | 214 ++++++++++++++++++++++ 1 file changed, 214 insertions(+) create mode 100644 scripts/convert_diffusers_to_cogvideox.py diff --git a/scripts/convert_diffusers_to_cogvideox.py b/scripts/convert_diffusers_to_cogvideox.py new file mode 100644 index 000000000000..3fa5d126bf19 --- /dev/null +++ b/scripts/convert_diffusers_to_cogvideox.py @@ -0,0 +1,214 @@ +import argparse +from typing import Any, Dict + +import torch + + +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: + if "state_dict" in saved_dict: + return saved_dict["state_dict"] + + if "model" in saved_dict: + return saved_dict["model"] + + return saved_dict + + +def rename_layernorm_key(key: str) -> str: + if ".attn1.norm_q." in key: + return key.replace( + ".attn1.norm_q.", + ".attention.query_layernorm_list.0.", + ) + + if ".attn1.norm_k." in key: + return key.replace( + ".attn1.norm_k.", + ".attention.key_layernorm_list.0.", + ) + + return key + + +def rename_transformer_key(key: str) -> str: + replacements = [ + ("transformer_blocks", "transformer.layers"), + (".attn1.to_out.0.", ".attention.dense."), + (".ff.net.0.proj.", ".mlp.dense_h_to_4h."), + (".ff.net.2.", ".mlp.dense_4h_to_h."), + ("norm_final.", "transformer.final_layernorm."), + ("time_embedding.linear_1.", "time_embed.0."), + ("time_embedding.linear_2.", "time_embed.2."), + ("ofs_embedding.linear_1.", "ofs_embed.0."), + ("ofs_embedding.linear_2.", "ofs_embed.2."), + ("patch_embed.", "mixins.patch_embed."), + ("norm_out.norm.", "mixins.final_layer.norm_final."), + ("norm_out.linear.", "mixins.final_layer.adaLN_modulation.1."), + ("proj_out.", "mixins.final_layer.linear."), + ] + + for old, new in replacements: + key = key.replace(old, new) + + return key + + +def convert_transformer_state_dict( + state_dict: Dict[str, Any], +) -> Dict[str, Any]: + converted_state_dict = {} + + # Reconstruct QKV. + for key in state_dict: + if key.endswith(("attn1.to_q.weight", "attn1.to_q.bias")): + to_k_key = key.replace("to_q", "to_k") + to_v_key = key.replace("to_q", "to_v") + + to_q = state_dict[key] + to_k = state_dict[to_k_key] + to_v = state_dict[to_v_key] + + new_key = rename_transformer_key( + key.replace( + ".attn1.to_q.", + ".attention.query_key_value.", + ) + ) + + converted_state_dict[new_key] = torch.cat( + [to_q, to_k, to_v], + dim=0, + ) + + # Copy ordinary keys. + qkv_parts = ( + ".attn1.to_q.", + ".attn1.to_k.", + ".attn1.to_v.", + ) + + for key, value in state_dict.items(): + if any(part in key for part in qkv_parts): + continue + + if ".attn1.norm_q." in key or ".attn1.norm_k." in key: + new_key = rename_layernorm_key(key) + new_key = rename_transformer_key(new_key) + converted_state_dict[new_key] = value + continue + + if ".norm1.linear." in key or ".norm2.linear." in key: + continue + + new_key = rename_transformer_key(key) + converted_state_dict[new_key] = value + + # Reconstruct AdaLN. + for layer_id in _get_layer_ids(state_dict): + norm1_prefix = ( + f"transformer_blocks.{layer_id}.norm1.linear." + ) + norm2_prefix = ( + f"transformer_blocks.{layer_id}.norm2.linear." + ) + + for weight_or_bias in ("weight", "bias"): + norm1 = state_dict[ + norm1_prefix + weight_or_bias + ] + norm2 = state_dict[ + norm2_prefix + weight_or_bias + ] + + chunks1 = norm1.chunk(6, dim=0) + chunks2 = norm2.chunk(6, dim=0) + + modulation = torch.cat( + [ + chunks1[0], + chunks1[1], + chunks1[2], + chunks2[0], + chunks2[1], + chunks2[2], + chunks1[3], + chunks1[4], + chunks1[5], + chunks2[3], + chunks2[4], + chunks2[5], + ], + dim=0, + ) + + converted_state_dict[ + "transformer.layers." + f"{layer_id}.adaln_layer." + f"adaLN_modulations.{weight_or_bias}" + ] = modulation + + return converted_state_dict + + +def _get_layer_ids( + state_dict: Dict[str, Any], +) -> list[str]: + layer_ids = set() + prefix = "transformer_blocks." + + for key in state_dict: + if key.startswith(prefix): + layer_ids.add( + key[len(prefix):].split(".", 1)[0] + ) + + return sorted(layer_ids, key=int) +def load_checkpoint(input_path: str) -> Dict[str, Any]: + print(f"Loading checkpoint from: {input_path}") + + checkpoint = torch.load( + input_path, + map_location="cpu", + ) + + state_dict = get_state_dict(checkpoint) + + print(f"Loaded {len(state_dict)} tensors") + + return state_dict +def get_args(): + parser = argparse.ArgumentParser( + description="Convert a Diffusers CogVideoX transformer checkpoint." + ) + + parser.add_argument( + "--input_path", + type=str, + required=True, + help="Path to the input Diffusers checkpoint.", + ) + + parser.add_argument( + "--output_path", + type=str, + required=True, + help="Path to save the converted checkpoint.", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + + print(f"Input: {args.input_path}") + print(f"Output: {args.output_path}") + + state_dict = load_checkpoint(args.input_path) + + converted = convert_transformer_state_dict(state_dict) + + print(f"Converted {len(converted)} tensors") + + for key, value in converted.items(): + print(key, tuple(value.shape)) \ No newline at end of file