Upload 5 files

app.py

@@ -11,7 +11,8 @@ import random
 import time
 
 from mod import (models, clear_cache, get_repo_safetensors, change_base_model,
-                 description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras, get_trigger_word, pipe)
+                 description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras,
+                 get_trigger_word, pipe, enhance_prompt)
 from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json,
                   download_my_lora, get_all_lora_tupled_list, apply_lora_prompt,
                   update_loras)
@@ -241,6 +242,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                         tagger_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use CogFlorence-2.1-Large", "Use Florence-2-Flux"], label="Algorithms", value=["Use WD Tagger"])
                         tagger_generate_from_image = gr.Button(value="Generate Prompt from Image")
                     prompt = gr.Textbox(label="Prompt", lines=1, max_lines=8, placeholder="Type a prompt")
+                    prompt_enhance = gr.Button(value="Enhance your prompt", variant="secondary")
             with gr.Column(scale=1, elem_id="gen_column"):
                 generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
         with gr.Row():
@@ -306,8 +308,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                             with gr.Accordion("From URL", open=True, visible=True):
                                 with gr.Row():
                                     lora_search_civitai_query = gr.Textbox(label="Query", placeholder="flux", lines=1)
-                                    lora_search_civitai_basemodel = gr.CheckboxGroup(label="Search LoRA for", choices=["Flux.1 D", "Flux.1 S"], value=["Flux.1 D", "Flux.1 S"])
                                     lora_search_civitai_submit = gr.Button("Search on Civitai")
+                                    lora_search_civitai_basemodel = gr.CheckboxGroup(label="Search LoRA for", choices=["Flux.1 D", "Flux.1 S"], value=["Flux.1 D", "Flux.1 S"])
                                 lora_search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value="", allow_custom_value=True, visible=False)
                                 lora_search_civitai_json = gr.JSON(value={}, visible=False)
                                 lora_search_civitai_desc = gr.Markdown(value="", visible=False)
@@ -344,6 +346,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
     )
 
     model_name.change(change_base_model, [model_name], [result])
+    prompt_enhance.click(enhance_prompt, [prompt], [prompt])
 
     gr.on(
         triggers=[lora_search_civitai_submit.click, lora_search_civitai_query.submit],

mod.py

@@ -7,6 +7,7 @@ import gc
 import subprocess
 
 
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 subprocess.run('pip cache purge', shell=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch.set_grad_enabled(False)
@@ -61,7 +62,7 @@ def get_repo_safetensors(repo_id: str):
         if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[])
         files = api.list_repo_files(repo_id=repo_id)
     except Exception as e:
-        print(f"Error: Failed to get {repo_id}'s info. ")
+        print(f"Error: Failed to get {repo_id}'s info.")
         print(e)
         return gr.update(choices=[])
     files = [f for f in files if f.endswith(".safetensors")]
@@ -138,8 +139,7 @@ def fuse_loras(pipe, lorajson: list[dict]):
     #pipe.unload_lora_weights()
 
 
-change_base_model.zerogpu = True
-fuse_loras.zerogpu = True
+
 
 
 def description_ui():
@@ -148,4 +148,31 @@ def description_ui():
 - Mod of [multimodalart/flux-lora-the-explorer](https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer),
  [gokaygokay/FLUX-Prompt-Generator](https://huggingface.co/spaces/gokaygokay/FLUX-Prompt-Generator).
 """
-    )
+    )
+
+
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+def load_prompt_enhancer():
+    try:
+        model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
+        tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).eval().to(device=device)
+        enhancer_flux = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device=device)
+    except Exception as e:
+        print(e)
+        enhancer_flux = None
+    return enhancer_flux
+
+
+enhancer_flux = load_prompt_enhancer()
+
+
+def enhance_prompt(input_prompt):
+    result = enhancer_flux("enhance prompt: " + input_prompt, max_length = 256)
+    enhanced_text = result[0]['generated_text']
+    return enhanced_text
+
+
+load_prompt_enhancer.zerogpu = True
+change_base_model.zerogpu = True
+fuse_loras.zerogpu = True

requirements.txt

@@ -1,7 +1,12 @@
 torch
+torchvision
+huggingface_hub
+accelerate
 git+https://github.com/huggingface/diffusers
 spaces
 transformers
 peft
 sentencepiece
-timm
+timm
+xformers
+einops