File size: 7,592 Bytes
699a234
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
import argparse
import json
import os
import shutil
from tempfile import TemporaryDirectory
from typing import List, Optional

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"]:
    config_file = "unet/config.json"
    os.makedirs(os.path.join(folder, "unet"), exist_ok=True)
    model_index_file = hf_hub_download(repo_id=model_id, filename="model_index.json")

    with open(model_index_file, "r") as f:
        index_dict = json.load(f)
        if not is_index_stable_diffusion_like(index_dict):
            print(f"{model_id} is not of type stable diffusion.")
            return False, False

    old_config_file = hf_hub_download(repo_id=model_id, filename=config_file)

    new_config_file = os.path.join(folder, config_file)
    success = convert_file(old_config_file, new_config_file)
    if success:
        operations = [CommitOperationAdd(path_in_repo=config_file, path_or_fileobj=new_config_file)]
        model_type = success
        return operations, model_type
    else:
        return False, False


def convert_file(
    old_config: str,
    new_config: str,
):
    with open(old_config, "r") as f:
        old_dict = json.load(f)

    is_stable_diffusion = "down_block_types" in old_dict and list(old_dict["down_block_types"]) == ["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"]

    is_stable_diffusion_1 = is_stable_diffusion and ("use_linear_projection" not in old_dict or old_dict["use_linear_projection"] is False)
    is_stable_diffusion_2 = is_stable_diffusion and ("use_linear_projection" in old_dict and old_dict["use_linear_projection"] is True)

    if not is_stable_diffusion_1 and not is_stable_diffusion_2:
        print("No matching config")
        return False

    if is_stable_diffusion_1:
        if old_dict["sample_size"] == 64:
            print("Dict correct")
            return False

        print("Correct stable diffusion 1")
        old_dict["sample_size"] = 64

    if is_stable_diffusion_2:
        if old_dict["sample_size"] == 96:
            print("Dict correct")
            return False

        print("Correct stable diffusion 2")
        old_dict["sample_size"] = 96

    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 1" if is_stable_diffusion_1 else "Stable Diffusion 2"


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 = "Correct `sample_size` of {}'s unet to have correct width and height default"
    info = api.model_info(model_id)
    filenames = set(s.rfilename for s in info.siblings)

    if "unet/config.json" not in filenames:
        print(f"Model: {model_id} has no 'unet/config.json' file to change")
        return

    if "vae/config.json" not in filenames:
        print(f"Model: {model_id} has no 'vae/config.json' file to change")
        return

    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, model_type = convert_single(model_id, folder)

            if operations:
                pr_title = pr_title.format(model_type)
                if model_type == "Stable Diffusion 1":
                    sample_size = 64
                    image_size = 512
                elif model_type == "Stable Diffusion 2":
                    sample_size = 96
                    image_size = 768

                pr_description = (
                        f"Since `diffusers==0.9.0` the width and height is automatically inferred from the `sample_size` attribute of your unet's config. It seems like your diffusion model has the same architecture as {model_type} which means that when using this model, by default an image size of {image_size}x{image_size} should be generated. This in turn means the unet's sample size should be **{sample_size}**. \n\n In order to suppress to update your configuration on the fly and to suppress the deprecation warning added in this PR: https://github.com/huggingface/diffusers/pull/1406/files#r1035703505 it is strongly recommended to merge this PR."
                )
                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)