Spaces:
Running
Running
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) |