Upload 2 files

app.py

@@ -66,7 +66,7 @@ def main():
                 input_tags_to_copy = gr.Textbox(value="", visible=False)
                 copy_input_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
                 translate_input_prompt_button = gr.Button(value="Translate prompt to English", size="sm", variant="secondary")
-                prompt_enhancer_model = gr.Radio(["Medium", "Long"], label="Model Choice", value="Long", info="Enhance your prompts with Medium or Long answers")
+                prompt_enhancer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Long", info="Enhance your prompts with Medium or Long answers")
                 with gr.Accordion(label="Advanced options", open=False, visible=False):
                     tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False)
                     dummy_np = gr.Textbox(label="Negative prompt", value="", visible=False)

promptenhancer.py

@@ -8,12 +8,12 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_models():
     try:
-        enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device="cpu")
-        enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device="cpu")
+        enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
+        enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
         model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
         tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).eval().to(device="cpu")
-        enhancer_flux = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device="cpu")
+        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_medium = enhancer_long = enhancer_flux = None
@@ -24,9 +24,7 @@ enhancer_medium, enhancer_long, enhancer_flux = load_models()
 @spaces.GPU
 def enhance_prompt(input_prompt, model_choice):
     if model_choice == "Medium":
-        enhancer_medium.to(device=device)
         result = enhancer_medium("Enhance the description: " + input_prompt)
-        enhancer_medium.to(device="cpu")
         enhanced_text = result[0]['summary_text']
         
         pattern = r'^.*?of\s+(.*?(?:\.|$))'
@@ -37,14 +35,10 @@ def enhance_prompt(input_prompt, model_choice):
             modified_sentence = match.group(1).capitalize()
             enhanced_text = modified_sentence + ' ' + remaining_text
     elif model_choice == "Flux":
-        enhancer_flux.to(device=device)
         result = enhancer_flux("enhance prompt: " + input_prompt, max_length = 256)
-        enhancer_flux.to(device="cpu")
         enhanced_text = result[0]['generated_text']
     else:  # Long
-        enhancer_long.to(device=device)
         result = enhancer_long("Enhance the description: " + input_prompt)
-        enhancer_long.to(device="cpu")
         enhanced_text = result[0]['summary_text']
     
     return enhanced_text