Function convert_controlnet

scripts/convert_cosmos_to_diffusers.py:642–715

Convert controlnet weights. Args: transformer_type: The type of transformer/controlnet control_state_dict: State dict containing controlnet-specific weights base_state_dict: State dict containing base transformer weights (for shared modules) weights_only: Wh

(
    transformer_type: str,
    control_state_dict: Dict[str, Any],
    base_state_dict: Dict[str, Any],
    weights_only: bool = True,
)

def convert_controlnet(
    transformer_type: str,
    control_state_dict: Dict[str, Any],
    base_state_dict: Dict[str, Any],
    weights_only: bool = True,
):
    """
    Convert controlnet weights.

    Args:
        transformer_type: The type of transformer/controlnet
        control_state_dict: State dict containing controlnet-specific weights
        base_state_dict: State dict containing base transformer weights (for shared modules)
        weights_only: Whether to use weights_only loading
    """
    if transformer_type not in CONTROLNET_CONFIGS:
        raise AssertionError(f"{transformer_type} does not define a ControlNet config")

    PREFIX_KEY = "net."

    # Process control-specific keys
    for key in list(control_state_dict.keys()):
        new_key = key[:]
        if new_key.startswith(PREFIX_KEY):
            new_key = new_key.removeprefix(PREFIX_KEY)
        for replace_key, rename_key in CONTROLNET_KEYS_RENAME_DICT.items():
            new_key = new_key.replace(replace_key, rename_key)
        update_state_dict_(control_state_dict, key, new_key)

    for key in list(control_state_dict.keys()):
        for special_key, handler_fn_inplace in CONTROLNET_SPECIAL_KEYS_REMAP.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, control_state_dict)

    # Copy shared weights from base transformer to controlnet
    # These are the duplicated modules: patch_embed_base, time_embed, learnable_pos_embed, img_context_proj, crossattn_proj
    shared_module_mappings = {
        # transformer key prefix -> controlnet key prefix
        "patch_embed.": "patch_embed_base.",
        "time_embed.": "time_embed.",
        "learnable_pos_embed.": "learnable_pos_embed.",
        "img_context_proj.": "img_context_proj.",
        "crossattn_proj.": "crossattn_proj.",
    }

    for key in list(base_state_dict.keys()):
        for transformer_prefix, controlnet_prefix in shared_module_mappings.items():
            if key.startswith(transformer_prefix):
                controlnet_key = controlnet_prefix + key[len(transformer_prefix) :]
                control_state_dict[controlnet_key] = base_state_dict[key].clone()
                print(f"Copied shared weight: {key} -> {controlnet_key}", flush=True)
                break

    cfg = CONTROLNET_CONFIGS[transformer_type]
    controlnet = CosmosControlNetModel(**cfg)

    expected_keys = set(controlnet.state_dict().keys())