Function convert_transformer

scripts/convert_skyreelsv2_to_diffusers.py:317–365

(model_type: str)

def convert_transformer(model_type: str):
    config = get_transformer_config(model_type)
    diffusers_config = config["diffusers_config"]
    model_id = config["model_id"]

    if "1.3B" in model_type:
        original_state_dict = load_file(hf_hub_download(model_id, "model.safetensors"))
    else:
        os.makedirs(model_type, exist_ok=True)
        model_dir = pathlib.Path(model_type)
        if "720P" in model_type:
            top_shard = 7 if "I2V" in model_type else 6
            zeros = "0" * (4 if "I2V" or "T2V" in model_type else 3)
            model_name = "diffusion_pytorch_model"
        elif "540P" in model_type:
            top_shard = 14 if "I2V" in model_type else 12
            model_name = "model"

        for i in range(1, top_shard + 1):
            shard_path = f"{model_name}-{i:05d}-of-{zeros}{top_shard}.safetensors"
            hf_hub_download(model_id, shard_path, local_dir=model_dir)
        original_state_dict = load_sharded_safetensors(model_dir)

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

    for key in list(original_state_dict.keys()):
        new_key = key[:]
        for replace_key, rename_key in TRANSFORMER_KEYS_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 TRANSFORMER_SPECIAL_KEYS_REMAP.items():
            if special_key not in key:
                continue
            handler_fn_inplace(key, original_state_dict)

    if "FLF2V" in model_type:
        if (
            hasattr(transformer.condition_embedder, "image_embedder")
            and hasattr(transformer.condition_embedder.image_embedder, "pos_embed")
            and transformer.condition_embedder.image_embedder.pos_embed is not None
        ):
            pos_embed_shape = transformer.condition_embedder.image_embedder.pos_embed.shape
            original_state_dict["condition_embedder.image_embedder.pos_embed"] = torch.zeros(pos_embed_shape)

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


def convert_vae():