File size: 5,996 Bytes
0163a2c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
import copy
import importlib
import os

import modules.scripts as scripts
from modules import paths, shared
from modules.processing import StableDiffusionProcessingImg2Img

original_alwayson_scripts = None


def find_controlnet():
    """Find ControlNet external_code

    Returns:
        module: ControlNet external_code module
    """
    try:
        cnet = importlib.import_module("extensions.sd-webui-controlnet.scripts.external_code")
    except Exception:
        try:
            cnet = importlib.import_module("extensions-builtin.sd-webui-controlnet.scripts.external_code")
        except Exception:
            cnet = None

    return cnet


def list_default_scripts():
    """Get list of default scripts

    Returns:
        list: List of default scripts
    """
    scripts_list = []

    basedir = os.path.join(paths.script_path, "scripts")
    if os.path.isdir(basedir):
        for filename in sorted(os.listdir(basedir)):
            if filename.endswith(".py"):
                scripts_list.append(filename)

    return scripts_list


def backup_alwayson_scripts(input_scripts):
    """Backup alwayson scripts

    Args:
        input_scripts (ScriptRunner): scripts to backup alwayson scripts
    """
    global original_alwayson_scripts
    original_alwayson_scripts = copy.copy(input_scripts.alwayson_scripts)


def disable_alwayson_scripts_wo_cn(cnet, input_scripts):
    """Disable alwayson scripts

    Args:
        input_scripts (ScriptRunner): scripts to disable alwayson scripts
    """
    default_scripts = list_default_scripts()

    disabled_scripts = []
    for script in input_scripts.alwayson_scripts:
        if os.path.basename(script.filename) in default_scripts:
            continue
        if cnet.is_cn_script(script):
            continue
        # print("Disabled script: {}".format(script.title()))
        disabled_scripts.append(script)

    for script in disabled_scripts:
        input_scripts.alwayson_scripts.remove(script)


def disable_all_alwayson_scripts(input_scripts):
    """Disable all alwayson scripts

    Args:
        input_scripts (ScriptRunner): scripts to disable alwayson scripts
    """
    default_scripts = list_default_scripts()

    disabled_scripts = []
    for script in input_scripts.alwayson_scripts:
        if os.path.basename(script.filename) in default_scripts:
            continue
        # print("Disabled script: {}".format(script.title()))
        disabled_scripts.append(script)

    for script in disabled_scripts:
        input_scripts.alwayson_scripts.remove(script)


def restore_alwayson_scripts(input_scripts):
    """Restore alwayson scripts

    Args:
        input_scripts (ScriptRunner): scripts to restore alwayson scripts
    """
    global original_alwayson_scripts
    if original_alwayson_scripts is not None:
        input_scripts.alwayson_scripts = original_alwayson_scripts
        original_alwayson_scripts = None


def get_max_args_to(input_scripts):
    """Get max args_to of scripts

    Args:
        input_scripts (ScriptRunner): scripts to get max args_to of scripts

    Returns:
        int: max args_to of scripts
    """
    max_args_to = 0
    for script in input_scripts.alwayson_scripts:
        if max_args_to < script.args_to:
            max_args_to = script.args_to
    return max_args_to


def get_controlnet_args_to(cnet, input_scripts):
    """Get args_to of ControlNet script

    Args:
        input_scripts (ScriptRunner): scripts to get args_to of ControlNet script

    Returns:
        int: args_to of ControlNet script
    """
    for script in input_scripts.alwayson_scripts:
        if cnet.is_cn_script(script):
            return script.args_to
    return get_max_args_to(input_scripts)


def clear_controlnet_cache(cnet, input_scripts):
    """Clear ControlNet cache

    Args:
        input_scripts (ScriptRunner): scripts to clear ControlNet cache
    """
    for script in input_scripts.alwayson_scripts:
        if cnet.is_cn_script(script):
            if hasattr(script, "clear_control_model_cache"):
                # print("Clear ControlNet cache: {}".format(script.title()))
                script.clear_control_model_cache()


def get_sd_img2img_processing(init_image, mask_image, prompt, n_prompt, sampler_id, ddim_steps, cfg_scale, strength, seed,
                              mask_blur=4, fill_mode=1):
    """Get StableDiffusionProcessingImg2Img instance

    Args:
        init_image (PIL.Image): Initial image
        mask_image (PIL.Image): Mask image
        prompt (str): Prompt
        n_prompt (int): Negative prompt
        sampler_id (str): Sampler ID
        ddim_steps (int): Steps
        cfg_scale (float): CFG scale
        strength (float): Denoising strength
        seed (int): Seed
        mask_blur (int, optional): Mask blur. Defaults to 4.
        fill_mode (int, optional): Fill mode. Defaults to 1.

    Returns:
        StableDiffusionProcessingImg2Img: StableDiffusionProcessingImg2Img instance
    """
    width, height = init_image.size

    sd_img2img_args = dict(
        sd_model=shared.sd_model,
        outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_img2img_samples,
        inpaint_full_res=False,
        init_images=[init_image],
        resize_mode=0,  # 0:Just resize
        denoising_strength=strength,
        image_cfg_scale=1.5,
        mask=mask_image,
        mask_blur=mask_blur,
        inpainting_fill=fill_mode,
        inpainting_mask_invert=0,   # 0:Inpaint masked
        prompt=prompt,
        negative_prompt=n_prompt,
        seed=seed,
        sampler_name=sampler_id,
        batch_size=1,
        n_iter=1,
        steps=ddim_steps,
        cfg_scale=cfg_scale,
        width=width,
        height=height,
        restore_faces=False,
        tiling=False,
        do_not_save_samples=True,
        do_not_save_grid=True,
    )

    p = StableDiffusionProcessingImg2Img(**sd_img2img_args)

    p.is_img2img = True
    p.scripts = scripts.scripts_img2img

    return p