File size: 12,813 Bytes
bbd4036
e530fe9
 
bbd4036
 
 
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e530fe9
 
 
 
 
 
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e530fe9
 
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e530fe9
 
 
 
920e1f5
e530fe9
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7bf31f
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7bf31f
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7bf31f
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e530fe9
 
920e1f5
 
 
 
 
 
 
 
 
 
 
bbd4036
 
920e1f5
 
 
 
 
e530fe9
 
 
 
 
 
 
 
920e1f5
 
 
 
 
 
 
 
 
c7bf31f
920e1f5
 
 
 
 
 
 
c7bf31f
 
920e1f5
 
 
 
 
 
 
 
 
 
 
 
 
 
e530fe9
920e1f5
e530fe9
920e1f5
 
 
 
e530fe9
 
 
920e1f5
 
 
c7bf31f
920e1f5
 
e530fe9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
920e1f5
 
e530fe9
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
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
import os
import time
import platform

def is_running_in_huggingface_spaces():
    return "SPACE_ID" in os.environ

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

try:
    # переводчик с русского на английский
    from google_translate import TranslatorWithCache
    is_google_translate_installed=True
    translator = TranslatorWithCache()
except ImportError:
    is_google_translate_installed=False


try:
    from config_ui import Config
    is_config_ui_installed=True
    config = Config()
    device = "cuda" if (config.cuda=="cuda" and torch.cuda.is_available()) else "cpu"
    lang=config.lang
except ImportError:
    is_config_ui_installed=False
    if platform.system() == "Darwin" and platform.machine().startswith("arm"):
        print("run on mac with Apple Silicon")
        if torch.backends.mps.is_available():
            device = torch.device("mps")  # MPS = Metal Performance Shaders
    else:
        #TODO parse env var to assign cuda device
        device = 0 if torch.cuda.is_available() else "cpu"
    lang='EN'

try:
    from prompt.portrait_prompt import generate_random_portrait_prompt
    is_rnd_gen_installed=True
except:
    is_rnd_gen_installed=False


model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).to(device)

max_target_length = 256
prefix = "enhance prompt"

def enhance_prompt(prompt, system_prompt, temperature=0.5, repetition_penalty=1.2, seed=-1, is_rnd_seed=True):
    global lang
    start_time = time.time()  # Начало замера времени
    

    if is_rnd_seed or seed==-1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
            
    torch.manual_seed(seed)
    
    if is_google_translate_installed:
        # Перевод с русского на английский
        en_prompt = translator.translate_ru2eng(prompt)
        input_text = f"{system_prompt}: {en_prompt}"
    else:
        input_text = f"{system_prompt}: {prompt}"
    
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    
    # Генерация текста
    outputs = model.generate(
        input_ids,
        max_length=max_target_length,
        num_return_sequences=1,
        do_sample=True,
        temperature=temperature,
        repetition_penalty=repetition_penalty
    )
    
    generated_text_en = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if is_google_translate_installed:
        result_output_ru = translator.translate_eng2ru(generated_text_en)
    else:
        result_output_ru=generated_text_en

    end_time = time.time()  # Конец замера времени
    execution_time = end_time - start_time
    time_str=f"execution time: {execution_time:.2f} s." if lang=="EN" else f"время выполнения: {execution_time:.2f} с."
    
    return seed, generated_text_en, result_output_ru, time_str

def random_prompt():
    rnd_prompt_str=generate_random_portrait_prompt()
    #rnd_prompt_str=get_random_words()
    return rnd_prompt_str

def copy_to_clipboard(text):
    gr.Info("скопировано в буффер обмена" if (lang=="RU") else "copy to clipboard" ,duration=1)
    return None

LABELS_EN={"prompt_input": "Input initial prompt:",
           "seed_output": "Seed:",
           "result_output" : "Improved prompt",
           "result_output_ru" : "Improved prompt (in Russian)",
           "generate_button": "Improve prompt",
           "copy_button": "Copy to clipboard",
           "save_button": "Save config",
           "advanced": "Advanced settings:",
           "system_prompt" : "System prompt:",
           "temperature": "Temperature",
           "repetition_penalty": "Repetition penalty",
           "is_rnd_seed": "Random Seed"
           }
LABELS=LABELS_EN

if is_google_translate_installed:

    LABELS_RU={"prompt_input": "Введите начальный промпт:",
            "seed_output": "Seed для генерации:",
            "result_output" : "Улучшенный промпт (на английском):",
            "result_output_ru" : "Улучшенный промпт (на русском):",
            "generate_button": "Улучшить промпт",
            "copy_button": "Скопировать в буффер обмена",
            "save_button": "Сохранить настройки",
            "advanced": "Расширенные настройки:",
            "system_prompt": "Системный промпт:",
            "temperature": "Температура",
            "repetition_penalty": "Штраф за повторение",
            "is_rnd_seed": "Случайный Seed"
            }
    LABELS=LABELS_EN if lang=="EN" else LABELS_RU

if is_google_translate_installed:
    def process_lang(selected_lang):
        global lang
        lang=selected_lang
        if selected_lang == "RU":
            LABELS=LABELS_RU
            message="Вы выбрали русский"
            isVisible=True
        elif selected_lang == "EN":
            LABELS=LABELS_EN
            message="You selected English"
            isVisible=False
        ret = [gr.update(value=LABELS["generate_button"]),
                gr.update(value=LABELS["copy_button"]),
                gr.update(value=LABELS["save_button"]),
                gr.update(label=LABELS["prompt_input"]),
                gr.update(label=LABELS["seed_output"]),
                gr.update(label=LABELS["is_rnd_seed"]),
                gr.update(label=LABELS["result_output"]),
                gr.update(visible=isVisible, label=LABELS["result_output_ru"]),
                gr.update(label=LABELS["advanced"]),
                gr.update(label=LABELS["system_prompt"]),
                gr.update(label=LABELS["temperature"]),
                gr.update(label=LABELS["repetition_penalty"])
            ]
        return message, *ret


if is_config_ui_installed:
    def save_config():
        global lang,device,isOpenAdvanced ,config, AccordionAdvanced

        config.set_lang(lang)
        config.set_cuda(str(device))
        isOpenAdvanced=AccordionAdvanced.open
        print(AccordionAdvanced.open)
        config.set_OpenAdvanced=(isOpenAdvanced)

        # Сохраняем изменения в файл
        config.save()
        return "save config to file" if lang=='EN' else "Конфигурация сохранена в файл"

def process_gpu(selected_gpu):
    """Функция для переключения модели между устройствами (CPU / CUDA)"""
    global model, device  # Используем глобальные переменные model и device
    device = torch.device(selected_gpu)  # Устанавливаем новое устройство
    model = model.to(device)  # Переносим модель на новое устройство
    message= f"Модель переключена на устройство: {selected_gpu}" if lang=="RU" else f"Model switched to device: {selected_gpu}"
    return message

def set_initial():
    global device
    dev="cpu"
    if str(device) =='cuda':
        device = torch.cuda.current_device()
        device_name = torch.cuda.get_device_name(device)
        device_name = f"GPU: {device_name}"
        dev="cuda"
    else:
        device_name = "running on CPU"
    return gr.update(value=lang), gr.update(value=dev), f'{device_name}, set to "{lang}" language'


# Настройка интерфейса Gradio
with gr.Blocks(title="Flux Prompt Enhance",
               theme=gr.themes.Default(primary_hue=gr.themes.colors.sky, secondary_hue=gr.themes.colors.indigo),
               analytics_enabled=False, css="footer{display:none !important}") as demo:
    gr.Image(label="header AiCave", value="./static/ai_cave_title.jpg",height="100%",
            show_download_button=False, show_label=False, show_share_button=False,
            interactive=False, show_fullscreen_button=False,)

    with gr.Row(variant="default"):
        running_in="for spaces" if is_running_in_huggingface_spaces() else "portable"
        gr.HTML(f'<h1>Flux Prompt Enhance {running_in} by <a href="https://boosty.to/aicave/donate" style="color: #4AA0E2;">CaveMan</a></h1>')
        with gr.Row(variant="default"):
        # выбор языка UI
            radio_lang = gr.Radio(choices = ["RU",  "EN"],  show_label = False, container = False, type = "value",
                                visible = True if is_google_translate_installed else False)

            if platform.system() == "Darwin" and platform.machine().startswith("arm") and torch.backends.mps.is_available():
                print("radio_gpu mac")
                radio_gpu  = gr.Radio(choices = ["mps","cpu"], show_label = False, container = False, type = "value",
                                    visible = True)
            else:
                radio_gpu  = gr.Radio(choices = ["cuda","cpu"], show_label = False, container = False, type = "value",
                                    visible = True if torch.cuda.is_available() else False)
            
            save_button = gr.Button(LABELS["save_button"], visible= True if is_config_ui_installed else False)
        
    with gr.Row(variant="default"):
        prompt_input     = gr.Textbox(label=LABELS["prompt_input"])
        if is_rnd_gen_installed:
            button_random = gr.Button("", icon="./static/random.png", scale=0,  min_width=200)
            button_random.click(fn=random_prompt, outputs=prompt_input)
    
    
    with gr.Accordion(label=LABELS["advanced"], open=False ) as AccordionAdvanced:
        with gr.Row(variant="default"):
            system_prompt = gr.Textbox(label=LABELS["system_prompt"], interactive=False,value=prefix)
            seed_output = gr.Textbox(label=LABELS["seed_output"], interactive=True,value=502119)
            is_rnd_seed = gr.Checkbox(value=True, label="Random seed", interactive=True)
            
        with gr.Row(variant="default"):           
            temperature = gr.Slider(label=LABELS["temperature"], interactive=True,value=0.7, minimum=0.1,maximum=1,step=0.1)
            repetition_penalty = gr.Slider(label=LABELS["repetition_penalty"], interactive=True,value=1.2, minimum=0.1,maximum=2,step=0.1)
            #repetition_penalty =

        
    result_output    = gr.Textbox(label=LABELS["result_output"], interactive=False)
    result_output_ru = gr.Textbox(label=LABELS["result_output_ru"], interactive=False, visible = False if lang == "EN" else True)
    
    #prompt_input.submit(fn=enhance_prompt, inputs=[prompt_input,system_prompt,temperature,repetition_penalty], outputs=[seed_output, result_output, result_output_ru], show_progress=False)

    # Кнопка генерации
    with gr.Row(variant="default"):
        generate_button = gr.Button(LABELS["generate_button"], variant="primary", size="lg")
    
    # Кнопка копирования в буфер обмена
        copy_button = gr.Button(LABELS["copy_button"], variant="secondary")
        copy_button.click(fn=copy_to_clipboard, inputs=result_output, outputs=[],js="(text) => navigator.clipboard.writeText(text)")
        
    with gr.Row(variant="default"):
        log_text = gr.Textbox(label="", container=False)
    
    if is_config_ui_installed:    
        save_button.click(fn=save_config, inputs=[], outputs=log_text)

    generate_button.click(fn=enhance_prompt, inputs=[prompt_input,system_prompt,temperature,repetition_penalty,seed_output,is_rnd_seed],
                              outputs=[seed_output, result_output, result_output_ru,log_text])
  
    if is_google_translate_installed:
        radio_lang.change(process_lang, inputs=radio_lang,
                        outputs=[log_text,generate_button, copy_button, save_button, prompt_input, seed_output, is_rnd_seed,
                                result_output, result_output_ru,AccordionAdvanced,system_prompt, temperature, repetition_penalty])

    radio_gpu.change(process_gpu, inputs=radio_gpu, outputs=log_text)
    
    #preload values for lang
    demo.load(set_initial, outputs=[radio_lang, radio_gpu, log_text])


launch_args={}

if not is_running_in_huggingface_spaces():
    launch_args["share"]=False
    launch_args["server_name"]="0.0.0.0"
    launch_args["inbrowser"] = True
    launch_args["favicon_path"] = "./static/favicon_aicave.png"
    launch_args["show_api"]=True
    
    if os.path.exists("cert.pem") and os.path.exists("key.pem"):
        launch_args["ssl_certfile"]="cert.pem"
        launch_args["ssl_keyfile"]="key.pem"
        launch_args["ssl_verify"]=False


demo.launch(**launch_args)