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)