| import argparse |
| import json |
| import os |
| import torch |
| import shutil |
| from tempfile import TemporaryDirectory |
| from typing import List, Optional |
| from diffusers import DiffusionPipeline |
|
|
| 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"]: |
| pipe = DiffusionPipeline.from_pretrained(model_id, cache_dir="/home/patrick/cache_to_delete") |
|
|
| try: |
| pipe.to(torch_dtype=torch.float16) |
| pipe.save_pretrained(folder, variant="fp16") |
| pipe.save_pretrained(folder, variant="fp16", safe_serialization=True) |
|
|
| all_files = [] |
| def find_files_in_dir(directory): |
| for root, dirs, files in os.walk(directory): |
| for file in files: |
| all_files.append(os.path.join(root, file)) |
|
|
| find_files_in_dir(folder) |
| files = [f for f in all_files if ".fp16." in f] |
|
|
| operations = [CommitOperationAdd(path_in_repo='/'.join(f.split("/")[-2:]), path_or_fileobj=f) for f in files] |
| return operations |
| except Exception as e: |
| print(e) |
| return False |
|
|
| def convert_file( |
| old_config: str, |
| new_config: str, |
| ): |
| with open(old_config, "r") as f: |
| old_dict = json.load(f) |
|
|
| old_dict["feature_extractor"][-1] = "CLIPImageProcessor" |
| |
| |
| |
| |
| |
| |
|
|
| 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" |
|
|
|
|
| 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 = "Fix deprecated float16/fp16 variant loading through new `version` API." |
|
|
| 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 = convert_single(model_id, folder) |
|
|
| if operations: |
| contributor = model_id.split("/")[0] |
| pr_description = ( |
| 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. " |
| "Loading `fp16` versions from a branch instead of the main branch is deprecated and will eventually be forbidden. " |
| "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." |
| 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:" |
| "\n\n" |
| "```py\n" |
| f"pipe = DiffusionPipeline.from_pretrained({model_id}, torch_dtype=torch.float16, variant='fp16')" |
| "\n```" |
| "\n\n" |
| "For more information please have a look at: https://huggingface.co/docs/diffusers/using-diffusers/loading#checkpoint-variants." |
| "\nWe made sure you that you can safely merge this pull request. \n\n Best, the 🧨 Diffusers team." |
| ) |
| 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) |
|
|
