File size: 1,601 Bytes
43c0fb7 08328de 43c0fb7 08328de 43c0fb7 08328de eb851d8 43c0fb7 eb851d8 08328de eb851d8 43c0fb7 08328de 43c0fb7 eb851d8 43c0fb7 9fa2bde eb851d8 9fa2bde 43c0fb7 eb851d8 43c0fb7 08328de 43c0fb7 08328de |
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 |
import os
import shutil
from tempfile import TemporaryDirectory
from typing import List, Optional, Tuple
import subprocess
from huggingface_hub import CommitOperationAdd, HfApi
from huggingface_hub.file_download import repo_folder_name
ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]]
def convert_to_core_ml(
model_id: str, folder: str, token: Optional[str], model_version: str, additional_args: str
) -> ConversionResult:
command = ["python3", "-m" , "python_coreml_stable_diffusion.torch2coreml", "--model-version", model_version, "-o", folder]
additional_args = additional_args
if additional_args == "":
# Set default args
additional_args = f"--convert-unet --convert-text-encoder --convert-vae-decoder --attention-implementation SPLIT_EINSUM --quantize-nbits 6"
command.extend(additional_args.split(" "))
print("Starting conversion: ", command)
subprocess.run(command)
print("Done")
api = HfApi(token=token)
api.upload_folder(
folder_path=folder,
repo_id=model_id,
path_in_repo="models",
repo_type="model",
)
def quantize(
api: "HfApi", model_id: str, model_version: str, additional_args: str
) -> None:
with TemporaryDirectory() as d:
folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
os.makedirs(folder)
try:
convert_to_core_ml(model_id, folder, token=api.token, model_version=model_version, additional_args=additional_args)
finally:
shutil.rmtree(folder) |