Function convert_transformer

scripts/convert_wan_to_diffusers.py:605–643

(model_type: str, stage: str = None)

def convert_transformer(model_type: str, stage: str = None):
    config, RENAME_DICT, SPECIAL_KEYS_REMAP = get_transformer_config(model_type)

    diffusers_config = config["diffusers_config"]
    model_id = config["model_id"]
    model_dir = pathlib.Path(snapshot_download(model_id, repo_type="model"))

    if stage is not None:
        model_dir = model_dir / stage

    original_state_dict = load_sharded_safetensors(model_dir)

    with init_empty_weights():
        if "Animate" in model_type:
            transformer = WanAnimateTransformer3DModel.from_config(diffusers_config)
        elif "VACE" in model_type:
            transformer = WanVACETransformer3DModel.from_config(diffusers_config)
        else:
            transformer = WanTransformer3DModel.from_config(diffusers_config)

    for key in list(original_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_(original_state_dict, key, new_key)

    for key in list(original_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, original_state_dict)

    # Load state dict into the meta model, which will materialize the tensors
    transformer.load_state_dict(original_state_dict, strict=True, assign=True)

    # Move to CPU to ensure all tensors are materialized
    transformer = transformer.to("cpu")

    return transformer


def convert_vae():