Update app.py

app.py

@@ -6,7 +6,7 @@ import random
 import os
 from PIL import Image
 from deep_translator import GoogleTranslator
-from gradio_client import Client  # Import the Gradio Client
+from gradio_client import Client
 
 # os.makedirs('assets', exist_ok=True)
 if not os.path.exists('icon.jpg'):
@@ -16,19 +16,15 @@ API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-
 timeout = 100
 
 def enhance_prompt(prompt):
-    """Enhance the prompt using the Mistral Nemo prompt enhancer API."""
-    client = Client("K00B404/mistral-nemo-prompt-enhancer")
+    client = Client("K00B404/Mistral-Nemo-custom")
     result = client.predict(
-        message=prompt,
-        system_message="You are an image generation prompt enhancer and should only respond with the enhanced version of the user input image generation prompt.",
-        max_tokens=512,
-        temperature=0.7,
-        top_p=0.95,
-        api_name="/chat"
+        system_prompt="You are a image generation prompt enhancer and must respond only with the enhanced version of the users input prompt",
+        user_message=f"###input image generation prompt### {prompt}",
+        api_name="/predict"
     )
     return result
 
-def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key=None, use_dev=False):
+def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key=None, use_dev=False, enhance_prompt_option=False):
     # Determine which API URL to use
     api_url = API_URL_DEV if use_dev else API_URL
 
@@ -53,11 +49,11 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
     print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
 
-    # Enhance the prompt using the API
-    enhanced_prompt = enhance_prompt(prompt)
-    print(f'\033[1mEnhanced Prompt:\033[0m {enhanced_prompt}')
+    if enhance_prompt_option:
+        prompt = enhance_prompt(prompt)
+        print(f'\033[1mGeneration {key} enhanced prompt:\033[0m {prompt}')
 
-    prompt = f"{enhanced_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'\033[1mGeneration {key}:\033[0m {prompt}')
 
     # If seed is -1, generate a random seed and use it
@@ -142,6 +138,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                         seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
                         huggingface_api_key = gr.Textbox(label="Hugging Face API Key (required for API calls)", placeholder="Enter your Hugging Face API Key here", type="password", elem_id="api-key")
                         use_dev = gr.Checkbox(label="Use Dev API", value=False, elem_id="use-dev-checkbox")
+                        enhance_prompt_option = gr.Checkbox(label="Enhance Prompt", value=False, elem_id="enhance-prompt-checkbox")
 
         with gr.Row():
             text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
@@ -149,7 +146,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
             image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
             seed_output = gr.Textbox(label="Seed Used", elem_id="seed-output")
         
-        # Adjust the click function to include the API key and use_dev as inputs
-        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev], outputs=[image_output, seed_output])
+        # Adjust the click function to include the API key, use_dev, and enhance_prompt_option as inputs
+        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev, enhance_prompt_option], outputs=[image_output, seed_output])
 
 app.launch(show_api=True, share=False,show_error=True)