File size: 6,777 Bytes
269cbe7 2acb9f2 269cbe7 2acb9f2 269cbe7 2acb9f2 269cbe7 2acb9f2 269cbe7 2acb9f2 269cbe7 2acb9f2 269cbe7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
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"
# if "clip_sample" not in old_dict:
# print("Make scheduler DDIM compatible")
# old_dict["clip_sample"] = False
# else:
# print("No matching config")
# return False
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)
|