Upload 3 files

app.py

@@ -0,0 +1,248 @@
+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
+    device = "cuda" 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):
+
+    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
+    
+    return seed, generated_text_en, result_output_ru
+
+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",
+           "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": "Сохранить настройки",
+            "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["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 = "use CPU"
+    return gr.update(value=lang), gr.update(value=dev), f'{device_name}\nset 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"):
+        gr.HTML("""
+        <h1>Flux Prompt Enhance portable 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)
+
+            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("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")
+        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])
+    
+    # Кнопка копирования в буфер обмена
+        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="")
+    
+    if is_config_ui_installed:    
+        save_button.click(fn=save_config, inputs=[], outputs=log_text)
+
+    #preload values for lang
+    demo.load(set_initial, outputs=[radio_lang, radio_gpu, 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,system_prompt, temperature, repetition_penalty])
+
+    radio_gpu.change(process_gpu, inputs=radio_gpu, outputs=log_text)
+
+
+# Запуск приложения с прослушиванием на всех интерфейсах и открытием в браузере
+demo.launch()

google_translate.py

@@ -0,0 +1,56 @@
+import re
+import time
+import random
+import os
+from googletrans import Translator
+from functools import lru_cache
+
+TRANSLATOR_CACHE_SIZE = int(os.getenv("TRANSLATOR_CACHE_SIZE",100))
+
+class TranslatorWithCache:
+    def __init__(self, cache_size=TRANSLATOR_CACHE_SIZE):
+        self.cache_size = cache_size
+        self._init_cache()
+
+    def _init_cache(self):
+        self.cached_translate = lru_cache(maxsize=self.cache_size)(self._translate)
+
+    def has_russian_word(self, text):
+        pattern = re.compile(r'[а-яА-ЯёЁ]+')
+        return bool(pattern.search(text))
+
+    def has_english_word(self, text):
+        pattern = re.compile(r'[a-zA-Z]+')
+        return bool(pattern.search(text))
+
+    def _translate(self, text: str, src: str, dest: str, retries=3) -> str:
+        tr = Translator()
+        for attempt in range(retries):
+            try:
+                return tr.translate(text, src=src, dest=dest).text
+            except TypeError as e:
+                if attempt < retries - 1:
+                    delay = random.uniform(1, 2)
+                    print(f"Ошибка: {e}. Повторная попытка через {delay:.2f} секунд...")
+                    time.sleep(delay)
+                else:
+                    print(f"Ошибка: {e}. Превышено количество попыток.")
+                    return text  # Возвращаем оригинальный текст, если все попытки провалились.
+
+    def translate_ru2eng(self, text: str, src='ru', dest='en') -> str:
+        if self.has_russian_word(text):
+            return self.cached_translate(text, src, dest)
+        else:
+            return text
+
+    def translate_eng2ru(self, text: str, src='en', dest='ru') -> str:
+        if self.has_english_word(text):
+            return self.cached_translate(text, src, dest)
+        else:
+            return text
+
+if __name__ == "__main__":
+    translator = TranslatorWithCache(cache_size=100)
+    print(translator.translate_ru2eng("Привет, как дела?"))
+    print(translator.translate_eng2ru("Hello, how are you?"))
+    print(translator.translate_ru2eng("Привет, как дела?"))  # Повторный вызов для проверки кэша

requirements.txt

@@ -0,0 +1,4 @@
+git+https://github.com/c-goosen/py-googletrans.git
+torch
+transformers
+gradio