diffusers/tools

1import argparse2import json3import os4import shutil5from tempfile import TemporaryDirectory6from typing import List, Optional7 8from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download9from huggingface_hub.file_download import repo_folder_name10 11 12class AlreadyExists(Exception):13    pass14 15 16def is_index_stable_diffusion_like(config_dict):17    if "_class_name" not in config_dict:18        return False19 20    compatible_classes = [21        "AltDiffusionImg2ImgPipeline",22        "AltDiffusionPipeline",23        "CycleDiffusionPipeline",24        "StableDiffusionImageVariationPipeline",25        "StableDiffusionImg2ImgPipeline",26        "StableDiffusionInpaintPipeline",27        "StableDiffusionInpaintPipelineLegacy",28        "StableDiffusionPipeline",29        "StableDiffusionPipelineSafe",30        "StableDiffusionUpscalePipeline",31        "VersatileDiffusionDualGuidedPipeline",32        "VersatileDiffusionImageVariationPipeline",33        "VersatileDiffusionPipeline",34        "VersatileDiffusionTextToImagePipeline",35        "OnnxStableDiffusionImg2ImgPipeline",36        "OnnxStableDiffusionInpaintPipeline",37        "OnnxStableDiffusionInpaintPipelineLegacy",38        "OnnxStableDiffusionPipeline",39        "StableDiffusionOnnxPipeline",40        "FlaxStableDiffusionPipeline",41    ]42    return config_dict["_class_name"] in compatible_classes43 44 45def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:46    config_file = "model_index.json"47    # os.makedirs(os.path.join(folder, "scheduler"), exist_ok=True)48    model_index_file = hf_hub_download(repo_id=model_id, filename="model_index.json")49 50    with open(model_index_file, "r") as f:51        index_dict = json.load(f)52        if index_dict.get("feature_extractor", None) is None:53            print(f"{model_id} has no feature extractor")54            return False, False55 56        if index_dict["feature_extractor"][-1] != "CLIPFeatureExtractor":57            print(f"{model_id} is not out of date or is not CLIP")58            return False, False59 60    # old_config_file = hf_hub_download(repo_id=model_id, filename=config_file)61    old_config_file = model_index_file62 63    new_config_file = os.path.join(folder, config_file)64    success = convert_file(old_config_file, new_config_file)65    if success:66        operations = [CommitOperationAdd(path_in_repo=config_file, path_or_fileobj=new_config_file)]67        model_type = success68        return operations, model_type69    else:70        return False, False71 72 73def convert_file(74    old_config: str,75    new_config: str,76):77    with open(old_config, "r") as f:78        old_dict = json.load(f)79 80    old_dict["feature_extractor"][-1] = "CLIPImageProcessor"81    # if "clip_sample" not in old_dict:82    #     print("Make scheduler DDIM compatible")83    #     old_dict["clip_sample"] = False84    # else:85    #     print("No matching config")86    #     return False87 88    with open(new_config, 'w') as f:89        json_str = json.dumps(old_dict, indent=2, sort_keys=True) + "\n"90        f.write(json_str)91 92    return "Stable Diffusion"93 94 95def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:96    try:97        discussions = api.get_repo_discussions(repo_id=model_id)98    except Exception:99        return None100    for discussion in discussions:101        if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:102            return discussion103 104 105def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:106#    pr_title = "Correct `sample_size` of {}'s unet to have correct width and height default"107    pr_title = "Fix deprecation warning by changing `CLIPFeatureExtractor` to `CLIPImageProcessor`."108    info = api.model_info(model_id)109    filenames = set(s.rfilename for s in info.siblings)110 111    if "model_index.json" not in filenames:112        print(f"Model: {model_id} has no model_index.json file to change")113        return114 115    # if "vae/config.json" not in filenames:116    #     print(f"Model: {model_id} has no 'vae/config.json' file to change")117    #     return118 119    with TemporaryDirectory() as d:120        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))121        os.makedirs(folder)122        new_pr = None123        try:124            operations = None125            pr = previous_pr(api, model_id, pr_title)126            if pr is not None and not force:127                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"128                new_pr = pr129                raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")130            else:131                operations, model_type = convert_single(model_id, folder)132 133            if operations:134                pr_title = pr_title.format(model_type)135#                if model_type == "Stable Diffusion 1":136#                    sample_size = 64137#                    image_size = 512138#                elif model_type == "Stable Diffusion 2":139#                    sample_size = 96140#                    image_size = 768141 142#                pr_description = (143#                        f"Since `diffusers==0.9.0` the width and height is automatically inferred from the `sample_size` attribute of your unet's config. It seems like your diffusion model has the same architecture as {model_type} which means that when using this model, by default an image size of {image_size}x{image_size} should be generated. This in turn means the unet's sample size should be **{sample_size}**. \n\n In order to suppress to update your configuration on the fly and to suppress the deprecation warning added in this PR: https://github.com/huggingface/diffusers/pull/1406/files#r1035703505 it is strongly recommended to merge this PR."144#                )145                contributor = model_id.split("/")[0]146                pr_description = (147                        f"Hey {contributor} 👋, \n\n Your model repository seems to contain logic to load a feature extractor that is deprecated, which you should notice by seeing the warning: "148                        "\n\n ```\ntransformers/models/clip/feature_extraction_clip.py:28: FutureWarning: The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. "149                        f"Please use CLIPImageProcessor instead. warnings.warn(\n``` \n\n when running `pipe = DiffusionPipeline.from_pretrained({model_id})`."150                        "This PR makes sure that the warning does not show anymore by replacing `CLIPFeatureExtractor` with `CLIPImageProcessor`. This will certainly not change or break your checkpoint, but only" 151                        "make sure that everything is up to date. \n\n Best, the 🧨 Diffusers team."152                )153                new_pr = api.create_commit(154                    repo_id=model_id,155                    operations=operations,156                    commit_message=pr_title,157                    commit_description=pr_description,158                    create_pr=True,159                )160                print(f"Pr created at {new_pr.pr_url}")161            else:162                print(f"No files to convert for {model_id}")163        finally:164            shutil.rmtree(folder)165        return new_pr166 167 168if __name__ == "__main__":169    DESCRIPTION = """170    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.171    It is PyTorch exclusive for now.172    It works by downloading the weights (PT), converting them locally, and uploading them back173    as a PR on the hub.174    """175    parser = argparse.ArgumentParser(description=DESCRIPTION)176    parser.add_argument(177        "model_id",178        type=str,179        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",180    )181    parser.add_argument(182        "--force",183        action="store_true",184        help="Create the PR even if it already exists of if the model was already converted.",185    )186    args = parser.parse_args()187    model_id = args.model_id188    api = HfApi()189    convert(api, model_id, force=args.force)190