|
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) |
|
|
