diffusers/tools

1import argparse2import json3import os4import shutil5from tempfile import TemporaryDirectory6from typing import List, Optional7 8from huggingface_hub import (9    CommitInfo,10    CommitOperationAdd,11    Discussion,12    HfApi,13    hf_hub_download,14)15from huggingface_hub.file_download import repo_folder_name16 17 18class AlreadyExists(Exception):19    pass20 21 22def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:23    config_file_name = "generation_config.json"24    config_file = hf_hub_download(repo_id=model_id, filename=config_file_name)25 26    old_config_file = config_file27 28    new_config_file = os.path.join(folder, config_file)29    success = convert_file(old_config_file, new_config_file)30    if success:31        operations = [32            CommitOperationAdd(33                path_in_repo=config_file_name, path_or_fileobj=new_config_file34            )35        ]36        model_type = success37        return operations, model_type38    else:39        return False, False40 41 42def convert_file(43    old_config: str,44    new_config: str,45):46    with open(old_config, "r") as f:47        old_dict = json.load(f)48 49    old_dict["max_initial_timestamp_index"] = 5050    old_dict["prev_sot_token_id"] = old_dict["suppress_tokens"][-2]51 52    with open(new_config, "w") as f:53        json_str = json.dumps(old_dict, indent=2, sort_keys=True) + "\n"54        f.write(json_str)55 56    return "Whisper"57 58 59def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:60    try:61        discussions = api.get_repo_discussions(repo_id=model_id)62    except Exception:63        return None64    for discussion in discussions:65        if (66            discussion.status == "open"67            and discussion.is_pull_request68            and discussion.title == pr_title69        ):70            return discussion71 72 73def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:74    pr_title = "Correct long-form generation config parameters 'max_initial_timestamp_index' and 'prev_sot_token_id'."75    info = api.model_info(model_id)76    filenames = set(s.rfilename for s in info.siblings)77 78    if "generation_config.json" not in filenames:79        print(f"Model: {model_id} has no generation_config.json file to change")80        return81 82    with TemporaryDirectory() as d:83        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))84        os.makedirs(folder)85        new_pr = None86        try:87            operations = None88            pr = previous_pr(api, model_id, pr_title)89            if pr is not None and not force:90                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"91                new_pr = pr92                raise AlreadyExists(93                    f"Model {model_id} already has an open PR check out {url}"94                )95            else:96                operations, model_type = convert_single(model_id, folder)97 98            if operations:99                pr_title = pr_title.format(model_type)100                contributor = model_id.split("/")[0]101                pr_description = (102                    f"Hey {contributor} 👋, \n\n Your model repository seems to contain outdated generation config parameters, such as 'max_initial_timestamp_index' and is missing the 'prev_sot_token_id' parameter. "103                    "These parameters need to be updated to correctly handle long-form generation as stated in  as part of https://github.com/huggingface/transformers/pull/27658. "104                    "This PR makes sure that everything is up to date and can be safely merged. \n\n Best, the Transformers team."105                )106                new_pr = api.create_commit(107                    repo_id=model_id,108                    operations=operations,109                    commit_message=pr_title,110                    commit_description=pr_description,111                    create_pr=True,112                )113                print(f"Pr created at {new_pr.pr_url}")114            else:115                print(f"No files to convert for {model_id}")116        finally:117            shutil.rmtree(folder)118        return new_pr119 120 121if __name__ == "__main__":122    DESCRIPTION = """123    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.124    It is PyTorch exclusive for now.125    It works by downloading the weights (PT), converting them locally, and uploading them back126    as a PR on the hub.127    """128    parser = argparse.ArgumentParser(description=DESCRIPTION)129    parser.add_argument(130        "model_id",131        type=str,132        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",133    )134    parser.add_argument(135        "--force",136        action="store_true",137        help="Create the PR even if it already exists of if the model was already converted.",138    )139    args = parser.parse_args()140    model_id = args.model_id141    api = HfApi()142    convert(api, model_id, force=args.force)143