| import argparse | |
| import json | |
| import os | |
| import shutil | |
| from tempfile import TemporaryDirectory | |
| from typing import List, Optional | |
| from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download | |
| from huggingface_hub.file_download import repo_folder_name | |
| class AlreadyExists(Exception): | |
| pass | |
| def is_index_stable_diffusion_like(config_dict): | |
| if "_class_name" not in config_dict: | |
| return False | |
| compatible_classes = [ | |
| "AltDiffusionImg2ImgPipeline", | |
| "AltDiffusionPipeline", | |
| "CycleDiffusionPipeline", | |
| "StableDiffusionImageVariationPipeline", | |
| "StableDiffusionImg2ImgPipeline", | |
| "StableDiffusionInpaintPipeline", | |
| "StableDiffusionInpaintPipelineLegacy", | |
| "StableDiffusionPipeline", | |
| "StableDiffusionPipelineSafe", | |
| "StableDiffusionUpscalePipeline", | |
| "VersatileDiffusionDualGuidedPipeline", | |
| "VersatileDiffusionImageVariationPipeline", | |
| "VersatileDiffusionPipeline", | |
| "VersatileDiffusionTextToImagePipeline", | |
| "OnnxStableDiffusionImg2ImgPipeline", | |
| "OnnxStableDiffusionInpaintPipeline", | |
| "OnnxStableDiffusionInpaintPipelineLegacy", | |
| "OnnxStableDiffusionPipeline", | |
| "StableDiffusionOnnxPipeline", | |
| "FlaxStableDiffusionPipeline", | |
| ] | |
| return config_dict["_class_name"] in compatible_classes | |
| def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]: | |
| config_file = "unet/config.json" | |
| os.makedirs(os.path.join(folder, "unet"), exist_ok=True) | |
| model_index_file = hf_hub_download(repo_id=model_id, filename="model_index.json") | |
| with open(model_index_file, "r") as f: | |
| index_dict = json.load(f) | |
| if not is_index_stable_diffusion_like(index_dict): | |
| print(f"{model_id} is not of type stable diffusion.") | |
| return False, False | |
| old_config_file = hf_hub_download(repo_id=model_id, filename=config_file) | |
| new_config_file = os.path.join(folder, config_file) | |
| success = convert_file(old_config_file, new_config_file) | |
| if success: | |
| operations = [CommitOperationAdd(path_in_repo=config_file, path_or_fileobj=new_config_file)] | |
| model_type = success | |
| return operations, model_type | |
| else: | |
| return False, False | |
| def convert_file( | |
| old_config: str, | |
| new_config: str, | |
| ): | |
| with open(old_config, "r") as f: | |
| old_dict = json.load(f) | |
| is_stable_diffusion = "down_block_types" in old_dict and list(old_dict["down_block_types"]) == ["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"] | |
| is_stable_diffusion_1 = is_stable_diffusion and ("use_linear_projection" not in old_dict or old_dict["use_linear_projection"] is False) | |
| is_stable_diffusion_2 = is_stable_diffusion and ("use_linear_projection" in old_dict and old_dict["use_linear_projection"] is True) | |
| if not is_stable_diffusion_1 and not is_stable_diffusion_2: | |
| print("No matching config") | |
| return False | |
| if is_stable_diffusion_1: | |
| if old_dict["sample_size"] == 64: | |
| print("Dict correct") | |
| return False | |
| print("Correct stable diffusion 1") | |
| old_dict["sample_size"] = 64 | |
| if is_stable_diffusion_2: | |
| if old_dict["sample_size"] == 96: | |
| print("Dict correct") | |
| return False | |
| print("Correct stable diffusion 2") | |
| old_dict["sample_size"] = 96 | |
| with open(new_config, 'w') as f: | |
| json_str = json.dumps(old_dict, indent=2, sort_keys=True) + "\n" | |
| f.write(json_str) | |
| return "Stable Diffusion 1" if is_stable_diffusion_1 else "Stable Diffusion 2" | |
| def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: | |
| try: | |
| discussions = api.get_repo_discussions(repo_id=model_id) | |
| except Exception: | |
| return None | |
| for discussion in discussions: | |
| if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title: | |
| return discussion | |
| def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]: | |
| pr_title = "Correct `sample_size` of {}'s unet to have correct width and height default" | |
| info = api.model_info(model_id) | |
| filenames = set(s.rfilename for s in info.siblings) | |
| if "unet/config.json" not in filenames: | |
| print(f"Model: {model_id} has no 'unet/config.json' file to change") | |
| return | |
| if "vae/config.json" not in filenames: | |
| print(f"Model: {model_id} has no 'vae/config.json' file to change") | |
| return | |
| with TemporaryDirectory() as d: | |
| folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) | |
| os.makedirs(folder) | |
| new_pr = None | |
| try: | |
| operations = None | |
| pr = previous_pr(api, model_id, pr_title) | |
| if pr is not None and not force: | |
| url = f"https://huggingface.co/{model_id}/discussions/{pr.num}" | |
| new_pr = pr | |
| raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}") | |
| else: | |
| operations, model_type = convert_single(model_id, folder) | |
| if operations: | |
| pr_title = pr_title.format(model_type) | |
| if model_type == "Stable Diffusion 1": | |
| sample_size = 64 | |
| image_size = 512 | |
| elif model_type == "Stable Diffusion 2": | |
| sample_size = 96 | |
| image_size = 768 | |
| pr_description = ( | |
| 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." | |
| ) | |
| new_pr = api.create_commit( | |
| repo_id=model_id, | |
| operations=operations, | |
| commit_message=pr_title, | |
| commit_description=pr_description, | |
| create_pr=True, | |
| ) | |
| print(f"Pr created at {new_pr.pr_url}") | |
| else: | |
| print(f"No files to convert for {model_id}") | |
| finally: | |
| shutil.rmtree(folder) | |
| return new_pr | |
| if __name__ == "__main__": | |
| DESCRIPTION = """ | |
| Simple utility tool to convert automatically some weights on the hub to `safetensors` format. | |
| It is PyTorch exclusive for now. | |
| It works by downloading the weights (PT), converting them locally, and uploading them back | |
| as a PR on the hub. | |
| """ | |
| parser = argparse.ArgumentParser(description=DESCRIPTION) | |
| parser.add_argument( | |
| "model_id", | |
| type=str, | |
| help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`", | |
| ) | |
| parser.add_argument( | |
| "--force", | |
| action="store_true", | |
| help="Create the PR even if it already exists of if the model was already converted.", | |
| ) | |
| args = parser.parse_args() | |
| model_id = args.model_id | |
| api = HfApi() | |
| convert(api, model_id, force=args.force) | |