Merge branch 'master' into cuda-device-id-selection

/SwinIR

/textual_inversion

.vscode

/extensions

- Estimated completion time in progress bar

- API

- Support for dedicated [inpainting model](https://github.com/runwayml/stable-diffusion#inpainting-with-stable-diffusion) by RunwayML. 

- Aesthetic Gradients, a way to generate images with a specific aesthetic by using clip images embds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients))

- via extension: [Aesthetic Gradients](https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients), a way to generate images with a specific aesthetic by using clip images embds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients))

## Where are Aesthetic Gradients?!?!

Aesthetic Gradients are now an extension. You can install it using git:

```commandline

git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients extensions/aesthetic-gradients

```

After running this command, make sure that you have `aesthetic-gradients` dir in webui's `extensions` directory and restart

the UI. The interface for Aesthetic Gradients should appear exactly the same as it was.

## Installation and Running

Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for both [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) and [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs.

    images_history_set_image_info(this); 

}

var images_history_click_tab = function(){

    var tabs_box = gradioApp().getElementById("images_history_tab");

    if (!tabs_box.classList.contains(this.getAttribute("tabname"))) {

        gradioApp().getElementById(this.getAttribute("tabname") + "_images_history_renew_page").click();

        tabs_box.classList.add(this.getAttribute("tabname"))

    }                

}

function images_history_disabled_del(){

    gradioApp().querySelectorAll(".images_history_del_button").forEach(function(btn){

        btn.setAttribute('disabled','disabled');

    var parent = item.parentElement;

    tagname = tagname.toUpperCase()

    while(parent.tagName != tagname){

        console.log(parent.tagName, tagname)

        parent = parent.parentElement;

    }  

    return parent;

}

function images_history_get_current_img(tabname, image_path, files){

function images_history_get_current_img(tabname, img_index, files){

    return [

        tabname, 

        gradioApp().getElementById(tabname + '_images_history_set_index').getAttribute("img_index"),         

        image_path, 

        files

    ];

}

function images_history_delete(del_num, tabname, img_path, img_file_name, page_index, filenames, image_index){

function images_history_delete(del_num, tabname, image_index){

    image_index = parseInt(image_index);

    var tab = gradioApp().getElementById(tabname + '_images_history');

    var set_btn = tab.querySelector(".images_history_set_index");

        }

    });    

    var img_num = buttons.length / 2;

    del_num = Math.min(img_num - image_index, del_num)    

    if (img_num <= del_num){

        setTimeout(function(tabname){

            gradioApp().getElementById(tabname + '_images_history_renew_page').click();

    } else {

        var next_img  

        for (var i = 0; i < del_num; i++){

            if (image_index + i < image_index + img_num){

            buttons[image_index + i].style.display = 'none';

                buttons[image_index + img_num + 1].style.display = 'none';

            buttons[image_index + i + img_num].style.display = 'none';

            next_img = image_index + i + 1

        }

        }

        var bnt;

        if (next_img  >= img_num){

            btn = buttons[image_index - del_num];

            btn = buttons[image_index - 1];

        } else {            

            btn = buttons[next_img];          

        } 

        setTimeout(function(btn){btn.click()}, 30, btn);

    }

    images_history_disabled_del();  

    return [del_num, tabname, img_path, img_file_name, page_index, filenames, image_index];

}

function images_history_turnpage(img_path, page_index, image_index, tabname){

function images_history_turnpage(tabname){

    gradioApp().getElementById(tabname + '_images_history_del_button').setAttribute('disabled','disabled');

    var buttons = gradioApp().getElementById(tabname + '_images_history').querySelectorAll(".gallery-item");

    buttons.forEach(function(elem) {

        elem.style.display = 'block';

    })   

    return [img_path, page_index, image_index, tabname];

}

function images_history_enable_del_buttons(){

}

function images_history_init(){ 

    var load_txt2img_button = gradioApp().getElementById('txt2img_images_history_renew_page')

    if (load_txt2img_button){        

    var tabnames = gradioApp().getElementById("images_history_tabnames_list")   

    if (tabnames){  

        images_history_tab_list = tabnames.querySelector("textarea").value.split(",")    

        for (var i in images_history_tab_list ){

            tab = images_history_tab_list[i];

            var tab = images_history_tab_list[i];

            gradioApp().getElementById(tab + '_images_history').classList.add("images_history_cantainor");

            gradioApp().getElementById(tab + '_images_history_set_index').classList.add("images_history_set_index");

            gradioApp().getElementById(tab + '_images_history_del_button').classList.add("images_history_del_button");

            gradioApp().getElementById(tab + '_images_history_gallery').classList.add("images_history_gallery");  

             gradioApp().getElementById(tab + "_images_history_start").setAttribute("style","padding:20px;font-size:25px");           

        }

        //preload

        if (gradioApp().getElementById("images_history_preload").querySelector("input").checked ){

            var tabs_box = gradioApp().getElementById("tab_images_history").querySelector("div").querySelector("div").querySelector("div");

            tabs_box.setAttribute("id", "images_history_tab");        

            var tab_btns = tabs_box.querySelectorAll("button"); 

            for (var i in images_history_tab_list){               

                var tabname = images_history_tab_list[i]

                tab_btns[i].setAttribute("tabname", tabname);

            // this refreshes history upon tab switch

            // until the history is known to work well, which is not the case now, we do not do this at startup

            //tab_btns[i].addEventListener('click', images_history_click_tab);

                tab_btns[i].addEventListener('click', function(){

                    var tabs_box = gradioApp().getElementById("images_history_tab");

                    if (!tabs_box.classList.contains(this.getAttribute("tabname"))) {

                        gradioApp().getElementById(this.getAttribute("tabname") + "_images_history_start").click();

                        tabs_box.classList.add(this.getAttribute("tabname"))

                    }                

                });

            }            

            tab_btns[0].click()

        }

        tabs_box.classList.add(images_history_tab_list[0]);

        // same as above, at page load

        //load_txt2img_button.click();

    } else {

        setTimeout(images_history_init, 500);

    } 

}

var images_history_tab_list = ["txt2img", "img2img", "extras"];

var images_history_tab_list = "";

setTimeout(images_history_init, 500);

document.addEventListener("DOMContentLoaded", function() {

    var mutationObserver = new MutationObserver(function(m){

        if (images_history_tab_list != ""){

            for (var i in images_history_tab_list ){

                let tabname = images_history_tab_list[i]

                var buttons = gradioApp().querySelectorAll('#' + tabname + '_images_history .gallery-item');

                    bnt.addEventListener('click', images_history_click_image, true);

                });

            // same as load_txt2img_button.click() above

            /*

                var cls_btn = gradioApp().getElementById(tabname + '_images_history_gallery').querySelector("svg");

                if (cls_btn){

                    cls_btn.addEventListener('click', function(){

                        gradioApp().getElementById(tabname + '_images_history_renew_page').click();

                    }, false);

            }*/

                }

            }     

        }

    });

    mutationObserver.observe(gradioApp(), { childList:true, subtree:true });

});

import copy

import itertools

import os

from pathlib import Path

import html

import gc

import gradio as gr

import torch

from PIL import Image

from torch import optim

from modules import shared

from transformers import CLIPModel, CLIPProcessor, CLIPTokenizer

from tqdm.auto import tqdm, trange

from modules.shared import opts, device

def get_all_images_in_folder(folder):

    return [os.path.join(folder, f) for f in os.listdir(folder) if

            os.path.isfile(os.path.join(folder, f)) and check_is_valid_image_file(f)]

def check_is_valid_image_file(filename):

    return filename.lower().endswith(('.png', '.jpg', '.jpeg', ".gif", ".tiff", ".webp"))

def batched(dataset, total, n=1):

    for ndx in range(0, total, n):

        yield [dataset.__getitem__(i) for i in range(ndx, min(ndx + n, total))]

def iter_to_batched(iterable, n=1):

    it = iter(iterable)

    while True:

        chunk = tuple(itertools.islice(it, n))

        if not chunk:

            return

        yield chunk

def create_ui():

    import modules.ui

    with gr.Group():

        with gr.Accordion("Open for Clip Aesthetic!", open=False):

            with gr.Row():

                aesthetic_weight = gr.Slider(minimum=0, maximum=1, step=0.01, label="Aesthetic weight",

                                             value=0.9)

                aesthetic_steps = gr.Slider(minimum=0, maximum=50, step=1, label="Aesthetic steps", value=5)

            with gr.Row():

                aesthetic_lr = gr.Textbox(label='Aesthetic learning rate',

                                          placeholder="Aesthetic learning rate", value="0.0001")

                aesthetic_slerp = gr.Checkbox(label="Slerp interpolation", value=False)

                aesthetic_imgs = gr.Dropdown(sorted(shared.aesthetic_embeddings.keys()),

                                             label="Aesthetic imgs embedding",

                                             value="None")

                modules.ui.create_refresh_button(aesthetic_imgs, shared.update_aesthetic_embeddings, lambda: {"choices": sorted(shared.aesthetic_embeddings.keys())}, "refresh_aesthetic_embeddings")

            with gr.Row():

                aesthetic_imgs_text = gr.Textbox(label='Aesthetic text for imgs',

                                                 placeholder="This text is used to rotate the feature space of the imgs embs",

                                                 value="")

                aesthetic_slerp_angle = gr.Slider(label='Slerp angle', minimum=0, maximum=1, step=0.01,

                                                  value=0.1)

                aesthetic_text_negative = gr.Checkbox(label="Is negative text", value=False)

    return aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative

aesthetic_clip_model = None

def aesthetic_clip():

    global aesthetic_clip_model

    if aesthetic_clip_model is None or aesthetic_clip_model.name_or_path != shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path:

        aesthetic_clip_model = CLIPModel.from_pretrained(shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path)

        aesthetic_clip_model.cpu()

    return aesthetic_clip_model

def generate_imgs_embd(name, folder, batch_size):

    model = aesthetic_clip().to(device)

    processor = CLIPProcessor.from_pretrained(model.name_or_path)

    with torch.no_grad():

        embs = []

        for paths in tqdm(iter_to_batched(get_all_images_in_folder(folder), batch_size),

                          desc=f"Generating embeddings for {name}"):

            if shared.state.interrupted:

                break

            inputs = processor(images=[Image.open(path) for path in paths], return_tensors="pt").to(device)

            outputs = model.get_image_features(**inputs).cpu()

            embs.append(torch.clone(outputs))

            inputs.to("cpu")

            del inputs, outputs

        embs = torch.cat(embs, dim=0).mean(dim=0, keepdim=True)

        # The generated embedding will be located here

        path = str(Path(shared.cmd_opts.aesthetic_embeddings_dir) / f"{name}.pt")

        torch.save(embs, path)

        model.cpu()

        del processor

        del embs

        gc.collect()

        torch.cuda.empty_cache()

        res = f"""

        Done generating embedding for {name}!

        Aesthetic embedding saved to {html.escape(path)}

        """

        shared.update_aesthetic_embeddings()

        return gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()), label="Imgs embedding",

                                  value="None"), \

               gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()),

                                  label="Imgs embedding",

                                  value="None"), res, ""

def slerp(low, high, val):

    low_norm = low / torch.norm(low, dim=1, keepdim=True)

    high_norm = high / torch.norm(high, dim=1, keepdim=True)

    omega = torch.acos((low_norm * high_norm).sum(1))

    so = torch.sin(omega)

    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high

    return res

class AestheticCLIP:

    def __init__(self):

        self.skip = False

        self.aesthetic_steps = 0

        self.aesthetic_weight = 0

        self.aesthetic_lr = 0

        self.slerp = False

        self.aesthetic_text_negative = ""

        self.aesthetic_slerp_angle = 0

        self.aesthetic_imgs_text = ""

        self.image_embs_name = None

        self.image_embs = None

        self.load_image_embs(None)

    def set_aesthetic_params(self, p, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, image_embs_name=None,

                             aesthetic_slerp=True, aesthetic_imgs_text="",

                             aesthetic_slerp_angle=0.15,

                             aesthetic_text_negative=False):

        self.aesthetic_imgs_text = aesthetic_imgs_text

        self.aesthetic_slerp_angle = aesthetic_slerp_angle

        self.aesthetic_text_negative = aesthetic_text_negative

        self.slerp = aesthetic_slerp

        self.aesthetic_lr = aesthetic_lr

        self.aesthetic_weight = aesthetic_weight

        self.aesthetic_steps = aesthetic_steps

        self.load_image_embs(image_embs_name)

        if self.image_embs_name is not None:

            p.extra_generation_params.update({

                "Aesthetic LR": aesthetic_lr,

                "Aesthetic weight": aesthetic_weight,

                "Aesthetic steps": aesthetic_steps,

                "Aesthetic embedding": self.image_embs_name,

                "Aesthetic slerp": aesthetic_slerp,

                "Aesthetic text": aesthetic_imgs_text,

                "Aesthetic text negative": aesthetic_text_negative,

                "Aesthetic slerp angle": aesthetic_slerp_angle,

            })

    def set_skip(self, skip):

        self.skip = skip

    def load_image_embs(self, image_embs_name):

        if image_embs_name is None or len(image_embs_name) == 0 or image_embs_name == "None":

            image_embs_name = None

            self.image_embs_name = None

        if image_embs_name is not None and self.image_embs_name != image_embs_name:

            self.image_embs_name = image_embs_name

            self.image_embs = torch.load(shared.aesthetic_embeddings[self.image_embs_name], map_location=device)

            self.image_embs /= self.image_embs.norm(dim=-1, keepdim=True)

            self.image_embs.requires_grad_(False)

    def __call__(self, z, remade_batch_tokens):

        if not self.skip and self.aesthetic_steps != 0 and self.aesthetic_lr != 0 and self.aesthetic_weight != 0 and self.image_embs_name is not None:

            tokenizer = shared.sd_model.cond_stage_model.tokenizer

            if not opts.use_old_emphasis_implementation:

                remade_batch_tokens = [

                    [tokenizer.bos_token_id] + x[:75] + [tokenizer.eos_token_id] for x in

                    remade_batch_tokens]

            tokens = torch.asarray(remade_batch_tokens).to(device)

            model = copy.deepcopy(aesthetic_clip()).to(device)

            model.requires_grad_(True)

            if self.aesthetic_imgs_text is not None and len(self.aesthetic_imgs_text) > 0:

                text_embs_2 = model.get_text_features(

                    **tokenizer([self.aesthetic_imgs_text], padding=True, return_tensors="pt").to(device))

                if self.aesthetic_text_negative:

                    text_embs_2 = self.image_embs - text_embs_2

                    text_embs_2 /= text_embs_2.norm(dim=-1, keepdim=True)

                img_embs = slerp(self.image_embs, text_embs_2, self.aesthetic_slerp_angle)

            else:

                img_embs = self.image_embs

            with torch.enable_grad():

                # We optimize the model to maximize the similarity

                optimizer = optim.Adam(

                    model.text_model.parameters(), lr=self.aesthetic_lr

                )

                for _ in trange(self.aesthetic_steps, desc="Aesthetic optimization"):

                    text_embs = model.get_text_features(input_ids=tokens)

                    text_embs = text_embs / text_embs.norm(dim=-1, keepdim=True)

                    sim = text_embs @ img_embs.T

                    loss = -sim

                    optimizer.zero_grad()

                    loss.mean().backward()

                    optimizer.step()

                zn = model.text_model(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)

                if opts.CLIP_stop_at_last_layers > 1:

                    zn = zn.hidden_states[-opts.CLIP_stop_at_last_layers]

                    zn = model.text_model.final_layer_norm(zn)

                else:

                    zn = zn.last_hidden_state

                model.cpu()

                del model

                gc.collect()

                torch.cuda.empty_cache()

            zn = torch.concat([zn[77 * i:77 * (i + 1)] for i in range(max(z.shape[1] // 77, 1))], 1)

            if self.slerp:

                z = slerp(z, zn, self.aesthetic_weight)

            else:

                z = z * (1 - self.aesthetic_weight) + zn * self.aesthetic_weight

        return z