Function convert_ltx2_transformer

scripts/convert_ltx2_to_diffusers.py:454–479

(original_state_dict: dict[str, Any], version: str)

def convert_ltx2_transformer(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
    config, rename_dict, special_keys_remap = get_ltx2_transformer_config(version)
    diffusers_config = config["diffusers_config"]

    transformer_state_dict, _ = split_transformer_and_connector_state_dict(original_state_dict)

    with init_empty_weights():
        transformer = LTX2VideoTransformer3DModel.from_config(diffusers_config)

    # Handle official code --> diffusers key remapping via the remap dict
    for key in list(transformer_state_dict.keys()):
        new_key = key[:]
        for replace_key, rename_key in rename_dict.items():
            new_key = new_key.replace(replace_key, rename_key)
        update_state_dict_inplace(transformer_state_dict, key, new_key)

    # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
    # special_keys_remap
    for key in list(transformer_state_dict.keys()):
        for special_key, handler_fn_inplace in special_keys_remap.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, transformer_state_dict)

    transformer.load_state_dict(transformer_state_dict, strict=True, assign=True)
    return transformer


def convert_ltx2_connectors(original_state_dict: dict[str, Any], version: str) -> LTX2TextConnectors: