diffusers/tools

1import argparse2import json3import os4import torch5import shutil6from tempfile import TemporaryDirectory7from typing import List, Optional8from diffusers import DiffusionPipeline9 10from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download11from huggingface_hub.file_download import repo_folder_name12 13 14class AlreadyExists(Exception):15    pass16 17 18def is_index_stable_diffusion_like(config_dict):19    if "_class_name" not in config_dict:20        return False21 22    compatible_classes = [23        "AltDiffusionImg2ImgPipeline",24        "AltDiffusionPipeline",25        "CycleDiffusionPipeline",26        "StableDiffusionImageVariationPipeline",27        "StableDiffusionImg2ImgPipeline",28        "StableDiffusionInpaintPipeline",29        "StableDiffusionInpaintPipelineLegacy",30        "StableDiffusionPipeline",31        "StableDiffusionPipelineSafe",32        "StableDiffusionUpscalePipeline",33        "VersatileDiffusionDualGuidedPipeline",34        "VersatileDiffusionImageVariationPipeline",35        "VersatileDiffusionPipeline",36        "VersatileDiffusionTextToImagePipeline",37        "OnnxStableDiffusionImg2ImgPipeline",38        "OnnxStableDiffusionInpaintPipeline",39        "OnnxStableDiffusionInpaintPipelineLegacy",40        "OnnxStableDiffusionPipeline",41        "StableDiffusionOnnxPipeline",42        "FlaxStableDiffusionPipeline",43    ]44    return config_dict["_class_name"] in compatible_classes45 46 47def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:48    pipe = DiffusionPipeline.from_pretrained(model_id, cache_dir="/home/patrick/cache_to_delete")49 50    try:51        pipe.to(torch_dtype=torch.float16)52        pipe.save_pretrained(folder, variant="fp16")53        pipe.save_pretrained(folder, variant="fp16", safe_serialization=True)54 55        all_files = []56        def find_files_in_dir(directory):57            for root, dirs, files in os.walk(directory):58                for file in files:59                    all_files.append(os.path.join(root, file))60 61        find_files_in_dir(folder)62        files = [f for f in all_files if ".fp16." in f]63 64        operations = [CommitOperationAdd(path_in_repo='/'.join(f.split("/")[-2:]), path_or_fileobj=f) for f in files]65        return operations66    except Exception as e:67        print(e)68        return False69 70def convert_file(71    old_config: str,72    new_config: str,73):74    with open(old_config, "r") as f:75        old_dict = json.load(f)76 77    old_dict["feature_extractor"][-1] = "CLIPImageProcessor"78    # if "clip_sample" not in old_dict:79    #     print("Make scheduler DDIM compatible")80    #     old_dict["clip_sample"] = False81    # else:82    #     print("No matching config")83    #     return False84 85    with open(new_config, 'w') as f:86        json_str = json.dumps(old_dict, indent=2, sort_keys=True) + "\n"87        f.write(json_str)88 89    return "Stable Diffusion"90 91 92def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:93    try:94        discussions = api.get_repo_discussions(repo_id=model_id)95    except Exception:96        return None97    for discussion in discussions:98        if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:99            return discussion100 101 102def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:103    pr_title = "Fix deprecated float16/fp16 variant loading through new `version` API."104 105    with TemporaryDirectory() as d:106        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))107        os.makedirs(folder)108        new_pr = None109        try:110            operations = None111            pr = previous_pr(api, model_id, pr_title)112            if pr is not None and not force:113                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"114                new_pr = pr115                raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")116            else:117                operations = convert_single(model_id, folder)118 119            if operations:120                contributor = model_id.split("/")[0]121                pr_description = (122                        f"Hey {contributor} 👋, \n\n Your model repository seems to contain a [`fp16` branch](https://huggingface.co/{model_id}/tree/fp16) to load the model in float16 precision. "123                        "Loading `fp16` versions from a branch instead of the main branch is deprecated and will eventually be forbidden. "124                        "Instead, we strongly recommend to save `fp16` versions of the model under `.fp16.` version files directly on the 'main' branch as enabled through this PR."125                        f"This PR makes sure that your model repository allows the user to correctly download float16 precision model weights by adding `fp16` model weights in both safetensors and PyTorch bin format:"126                        "\n\n"127                        "```py\n"128                        f"pipe = DiffusionPipeline.from_pretrained({model_id}, torch_dtype=torch.float16, variant='fp16')"129                        "\n```"130                        "\n\n"131                        "For more information please have a look at: https://huggingface.co/docs/diffusers/using-diffusers/loading#checkpoint-variants."132                        "\nWe made sure you that you can safely merge this pull request. \n\n Best, the 🧨 Diffusers team."133                )134                new_pr = api.create_commit(135                    repo_id=model_id,136                    operations=operations,137                    commit_message=pr_title,138                    commit_description=pr_description,139                    create_pr=True,140                )141                print(f"Pr created at {new_pr.pr_url}")142            else:143                print(f"No files to convert for {model_id}")144        finally:145            shutil.rmtree(folder)146        return new_pr147 148 149if __name__ == "__main__":150    DESCRIPTION = """151    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.152    It is PyTorch exclusive for now.153    It works by downloading the weights (PT), converting them locally, and uploading them back154    as a PR on the hub.155    """156    parser = argparse.ArgumentParser(description=DESCRIPTION)157    parser.add_argument(158        "model_id",159        type=str,160        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",161    )162    parser.add_argument(163        "--force",164        action="store_true",165        help="Create the PR even if it already exists of if the model was already converted.",166    )167    args = parser.parse_args()168    model_id = args.model_id169    api = HfApi()170    convert(api, model_id, force=args.force)171