aboutsummaryrefslogtreecommitdiff
path: root/modules/inspiration.py
blob: 319183ab41b97f074bc8165dabb296194cd48d16 (plain)
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
import os
import random
import gradio
from modules.shared import opts
inspiration_system_path = os.path.join(opts.inspiration_dir, "system")
def read_name_list(file, types=None, keyword=None):
    if not os.path.exists(file):
        return []
    ret = []
    f = open(file, "r")    
    line = f.readline()
    while len(line) > 0:
        line = line.rstrip("\n")
        if types is not None:
            dirname = os.path.split(line)
            if dirname[0] in types and keyword in dirname[1].lower():
                ret.append(line)
        else:
            ret.append(line)
        line = f.readline()
    return ret

def save_name_list(file, name): 
    name_list = read_name_list(file)
    if name not in name_list:
        with open(file, "a") as f:
            f.write(name + "\n")

def get_types_list():
    files = os.listdir(opts.inspiration_dir)
    types = []
    for x in files:
        path = os.path.join(opts.inspiration_dir, x)
        if x[0] == ".":
            continue
        if not os.path.isdir(path):
            continue
        if path == inspiration_system_path:
            continue
        types.append(x)
    return types

def get_inspiration_images(source, types, keyword):
    keyword = keyword.strip(" ").lower()
    get_num = int(opts.inspiration_rows_num * opts.inspiration_cols_num)
    if source == "Favorites":
        names = read_name_list(os.path.join(inspiration_system_path, "faverites.txt"), types, keyword)
        names = random.sample(names, get_num) if len(names) > get_num else names
    elif source == "Abandoned":
        names = read_name_list(os.path.join(inspiration_system_path, "abandoned.txt"), types, keyword)
        names = random.sample(names, get_num) if len(names) > get_num else names
    elif source == "Exclude abandoned":        
        abandoned = read_name_list(os.path.join(inspiration_system_path, "abandoned.txt"), types, keyword)  
        all_names = []
        for tp in types:
            name_list = os.listdir(os.path.join(opts.inspiration_dir, tp))
            all_names += [os.path.join(tp, x) for x in name_list if keyword in x.lower()]
        
        if len(all_names) > get_num:
            names = []
            while len(names) < get_num:
                name = random.choice(all_names)
                if name not in abandoned:
                    names.append(name)
        else:
            names = all_names
    else:
        all_names = []
        for tp in types:
            name_list = os.listdir(os.path.join(opts.inspiration_dir, tp))
            all_names += [os.path.join(tp, x) for x in name_list if keyword in x.lower()]
        names = random.sample(all_names, get_num) if len(all_names) > get_num else all_names
    image_list = []
    for a in names:
        image_path = os.path.join(opts.inspiration_dir, a)
        images = os.listdir(image_path)        
        image_list.append((os.path.join(image_path, random.choice(images)), a))
    return image_list, names

def select_click(index, name_list):
    name = name_list[int(index)]
    path = os.path.join(opts.inspiration_dir, name)
    images = os.listdir(path)
    return name, [os.path.join(path, x) for x in images], ""

def give_up_click(name):
    file = os.path.join(inspiration_system_path, "abandoned.txt")
    save_name_list(file, name)
    return "Added to abandoned list"
   
def collect_click(name):
    file = os.path.join(inspiration_system_path, "faverites.txt")
    save_name_list(file, name)
    return "Added to faverite list"

def moveout_click(name, source):
    if source == "Abandoned":
        file = os.path.join(inspiration_system_path, "abandoned.txt")
    elif source == "Favorites":
        file = os.path.join(inspiration_system_path, "faverites.txt")
    else:
        return None
    name_list = read_name_list(file)
    os.remove(file)
    with open(file, "a") as f:
        for a in name_list:
            if a != name:
                f.write(a + "\n")
    return f"Moved out {name} from {source} list"

def source_change(source):
    if source in ["Abandoned", "Favorites"]:
        return gradio.update(visible=True), []
    else:
        return gradio.update(visible=False), []
def add_to_prompt(name, prompt):
    name = os.path.basename(name)
    return prompt + "," + name

def clear_keyword():
    return ""

def ui(gr, opts, txt2img_prompt, img2img_prompt):     
    with gr.Blocks(analytics_enabled=False) as inspiration:
        flag = os.path.exists(opts.inspiration_dir)        
        if flag:
            types = get_types_list()
            flag = len(types) > 0
        else:
            os.makedirs(opts.inspiration_dir)
        if not flag:            
            gr.HTML("""
                <div align='center' width="50%"><h2>To activate inspiration function, you need get "inspiration" images first. </h2><br>
                You can create these images by run "Create inspiration images" script in txt2img page, <br> you can get the artists or art styles list from here<br>
                <a href="https://github.com/pharmapsychotic/clip-interrogator/tree/main/data">https://github.com/pharmapsychotic/clip-interrogator/tree/main/data</a><br>
                download these files, and select these files in the "Create inspiration images" script UI<br>
                There about 6000 artists and art styles in these files. <br>This takes server hours depending on your GPU type and how many pictures  you generate for each artist/style
                <br>I suggest at least four images for each<br><br><br>
                <h2>You can also download generated pictures from here:</h2><br>
                <a href="https://huggingface.co/datasets/yfszzx/inspiration">https://huggingface.co/datasets/yfszzx/inspiration</a><br>
                unzip the file to the project directory of webui<br>
                and restart webui, and enjoy the joy of creation!<br></div>
                """)      
            return inspiration
        if not os.path.exists(inspiration_system_path):
            os.mkdir(inspiration_system_path)
        with gr.Row():
            with gr.Column(scale=2):                
                inspiration_gallery = gr.Gallery(show_label=False, elem_id="inspiration_gallery").style(grid=opts.inspiration_cols_num, height='auto')
            with gr.Column(scale=1):
                types = gr.CheckboxGroup(choices=types, value=types)
                with gr.Row(): 
                    source = gr.Dropdown(choices=["All", "Favorites", "Exclude abandoned", "Abandoned"], value="Exclude abandoned", label="Source")
                    keyword = gr.Textbox("", label="Key word")                                                
                get_inspiration = gr.Button("Get inspiration", elem_id="inspiration_get_button")   
                name = gr.Textbox(show_label=False, interactive=False)      
                with gr.Row():                     
                    send_to_txt2img = gr.Button('to txt2img')
                    send_to_img2img = gr.Button('to img2img')
                style_gallery = gr.Gallery(show_label=False).style(grid=2, height='auto') 
                warning = gr.HTML()
                with gr.Row():
                    collect = gr.Button('Collect')     
                    give_up = gr.Button("Don't show again")
                moveout = gr.Button("Move out", visible=False)
                
        with gr.Row(visible=False):
            select_button = gr.Button('set button', elem_id="inspiration_select_button")
            name_list = gr.State()
        
        get_inspiration.click(get_inspiration_images, inputs=[source, types, keyword], outputs=[inspiration_gallery, name_list])
        keyword.submit(fn=None, _js="inspiration_click_get_button", inputs=None, outputs=None)
        source.change(source_change, inputs=[source], outputs=[moveout, style_gallery])
        source.change(fn=clear_keyword, _js="inspiration_click_get_button", inputs=None, outputs=[keyword])       
        types.change(fn=clear_keyword, _js="inspiration_click_get_button", inputs=None, outputs=[keyword])  

        select_button.click(select_click, _js="inspiration_selected", inputs=[name, name_list], outputs=[name, style_gallery, warning])
        give_up.click(give_up_click, inputs=[name], outputs=[warning])
        collect.click(collect_click, inputs=[name], outputs=[warning])
        moveout.click(moveout_click, inputs=[name, source], outputs=[warning])
        moveout.click(fn=None, _js="inspiration_click_get_button", inputs=None, outputs=None)

        send_to_txt2img.click(add_to_prompt, inputs=[name, txt2img_prompt], outputs=[txt2img_prompt])
        send_to_img2img.click(add_to_prompt, inputs=[name, img2img_prompt], outputs=[img2img_prompt])
        send_to_txt2img.click(collect_click, inputs=[name], outputs=[warning])
        send_to_img2img.click(collect_click, inputs=[name], outputs=[warning])
        send_to_txt2img.click(None, _js='switch_to_txt2img', inputs=None, outputs=None)
        send_to_img2img.click(None, _js="switch_to_img2img_img2img", inputs=None, outputs=None)
    return inspiration