Merge remote-tracking branch 'upstream/master' into ModelLoader

/tmp

/model.ckpt

/models/**/*

/GFPGANv1.3.pth

/gfpgan/weights/*.pth

/ui-config.json

/outputs

/config.json

    "\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomed",

    "\u{1f3a8}": "Add a random artist to the prompt.",

    "\u2199\ufe0f": "Read generation parameters from prompt into user interface.",

    "\uD83D\uDCC2": "Open images output directory",

    "Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt",

    "SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back",

    });

    json_elem.parentElement.style.display="none"

	if (!txt2img_textarea) {

		txt2img_textarea = gradioApp().querySelector("#txt2img_prompt > label > textarea");

		txt2img_textarea?.addEventListener("input", () => update_token_counter("txt2img_token_button"));

	}

	if (!img2img_textarea) {

		img2img_textarea = gradioApp().querySelector("#img2img_prompt > label > textarea");

		img2img_textarea?.addEventListener("input", () => update_token_counter("img2img_token_button"));

	}

})

let txt2img_textarea, img2img_textarea = undefined;

let wait_time = 800

let token_timeout;

function update_token_counter(button_id) {

	if (token_timeout)

		clearTimeout(token_timeout);

	token_timeout = setTimeout(() => gradioApp().getElementById(button_id)?.click(), wait_time);

}

requirements_file = os.environ.get('REQS_FILE', "requirements_versions.txt")

commandline_args = os.environ.get('COMMANDLINE_ARGS', "")

k_diffusion_package = os.environ.get('K_DIFFUSION_PACKAGE', "git+https://github.com/crowsonkb/k-diffusion.git@1a0703dfb7d24d8806267c3e7ccc4caf67fd1331")

gfpgan_package = os.environ.get('GFPGAN_PACKAGE', "git+https://github.com/TencentARC/GFPGAN.git@8d2447a2d918f8eba5a4a01463fd48e45126a379")

stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc")

taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")

k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "a7ec1974d4ccb394c2dca275f42cd97490618924")

codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")

blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")

    run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch")

if not skip_torch_cuda_test:

    run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDINE_ARGS variable to disable this check'")

if not is_installed("k_diffusion.sampling"):

    run_pip(f"install {k_diffusion_package}", "k-diffusion")

    run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check'")

if not is_installed("gfpgan"):

    run_pip(f"install {gfpgan_package}", "gfpgan")

git_clone("https://github.com/CompVis/stable-diffusion.git", repo_dir('stable-diffusion'), "Stable Diffusion", stable_diffusion_commit_hash)

git_clone("https://github.com/CompVis/taming-transformers.git", repo_dir('taming-transformers'), "Taming Transformers", taming_transformers_commit_hash)

git_clone("https://github.com/crowsonkb/k-diffusion.git", repo_dir('k-diffusion'), "K-diffusion", k_diffusion_commit_hash)

git_clone("https://github.com/sczhou/CodeFormer.git", repo_dir('CodeFormer'), "CodeFormer", codeformer_commit_hash)

git_clone("https://github.com/salesforce/BLIP.git", repo_dir('BLIP'), "BLIP", blip_commit_hash)

if os.path.isdir(repo_dir('latent-diffusion')):

sys.argv += args

if "--exit" in args:

    print("Exiting because of --exit argument")

    exit(0)

def start_webui():

    print(f"Launching Web UI with arguments: {' '.join(sys.argv[1:])}")

import torch

import tqdm

from modules import processing, shared, images, devices

from modules import processing, shared, images, devices, sd_models

from modules.shared import opts

import modules.gfpgan_model

from modules.ui import plaintext_to_html

import modules.codeformer_model

import piexif

import piexif.helper

import gradio as gr

cached_images = {}

    return '', geninfo, info

def run_modelmerger(modelname_0, modelname_1, interp_method, interp_amount):

def run_modelmerger(primary_model_name, secondary_model_name, interp_method, interp_amount, save_as_half, custom_name):

    # Linear interpolation (https://en.wikipedia.org/wiki/Linear_interpolation)

    def weighted_sum(theta0, theta1, alpha):

        return ((1 - alpha) * theta0) + (alpha * theta1)

        alpha = alpha * alpha * (3 - (2 * alpha))

        return theta0 + ((theta1 - theta0) * alpha)

    if os.path.exists(modelname_0):

        model0_filename = modelname_0

        modelname_0 = os.path.splitext(os.path.basename(modelname_0))[0]

    else:

        model0_filename = 'models/' + modelname_0 + '.ckpt'

    # Inverse Smoothstep (https://en.wikipedia.org/wiki/Smoothstep)

    def inv_sigmoid(theta0, theta1, alpha):

        import math

        alpha = 0.5 - math.sin(math.asin(1.0 - 2.0 * alpha) / 3.0)

        return theta0 + ((theta1 - theta0) * alpha)

    if os.path.exists(modelname_1):

        model1_filename = modelname_1

        modelname_1 = os.path.splitext(os.path.basename(modelname_1))[0]

    else:

        model1_filename = 'models/' + modelname_1 + '.ckpt'

    primary_model_info = sd_models.checkpoints_list[primary_model_name]

    secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

    print(f"Loading {model0_filename}...")

    model_0 = torch.load(model0_filename, map_location='cpu')

    print(f"Loading {primary_model_info.filename}...")

    primary_model = torch.load(primary_model_info.filename, map_location='cpu')

    print(f"Loading {model1_filename}...")

    model_1 = torch.load(model1_filename, map_location='cpu')

    print(f"Loading {secondary_model_info.filename}...")

    secondary_model = torch.load(secondary_model_info.filename, map_location='cpu')

    theta_0 = model_0['state_dict']

    theta_1 = model_1['state_dict']

    theta_0 = primary_model['state_dict']

    theta_1 = secondary_model['state_dict']

    theta_funcs = {

        "Weighted Sum": weighted_sum,

        "Sigmoid": sigmoid,

        "Inverse Sigmoid": inv_sigmoid,

    }

    theta_func = theta_funcs[interp_method]

    print(f"Merging...")

    for key in tqdm.tqdm(theta_0.keys()):

        if 'model' in key and key in theta_1:

            theta_0[key] = theta_func(theta_0[key], theta_1[key], interp_amount)

            theta_0[key] = theta_func(theta_0[key], theta_1[key], (float(1.0) - interp_amount))  # Need to reverse the interp_amount to match the desired mix ration in the merged checkpoint

            if save_as_half:

                theta_0[key] = theta_0[key].half()

    for key in theta_1.keys():

        if 'model' in key and key not in theta_0:

            theta_0[key] = theta_1[key]

            if save_as_half:

                theta_0[key] = theta_0[key].half()

    filename = primary_model_info.model_name + '_' + str(round(interp_amount, 2)) + '-' + secondary_model_info.model_name + '_' + str(round((float(1.0) - interp_amount), 2)) + '-' + interp_method.replace(" ", "_") + '-merged.ckpt'

    filename = filename if custom_name == '' else (custom_name + '.ckpt')

    output_modelname = os.path.join(shared.cmd_opts.ckpt_dir, filename)

    output_modelname = 'models/' + modelname_0 + '-' + modelname_1 + '-merged.ckpt'

    print(f"Saving to {output_modelname}...")

    torch.save(model_0, output_modelname)

    torch.save(primary_model, output_modelname)

    sd_models.list_models()

    print(f"Checkpoint saved.")

    return "Checkpoint saved to " + output_modelname

    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(3)]