Update app.py

app.py

@@ -19,8 +19,12 @@ timeout = 100
 # Initialize the prompt enhancer client
 prompt_enhancer = Client("K00B404/mistral-nemo-prompt-enhancer")
 
-def enhance_prompt(prompt):
+def enhance_prompt(prompt, enable_enhancement=True):
     """Enhance the given prompt using the Mistral Nemo prompt enhancer."""
+    if not enable_enhancement:
+        print(f'\033[1mPrompt enhancement disabled, using original prompt:\033[0m {prompt}')
+        return prompt
+        
     try:
         system_message = "You are an expert at writing detailed, high-quality image generation prompts. Enhance the given prompt by adding relevant artistic details, style elements, and quality descriptors. Keep the original intent but make it more elaborate and specific."
         enhanced = prompt_enhancer.predict(
@@ -39,9 +43,9 @@ def enhance_prompt(prompt):
         return prompt  # Fall back to original prompt if enhancement fails
 
 # Function to query the API and return the generated image
-def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
+def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, enable_enhancement=True):
     if prompt == "" or prompt is None:
-        return None
+        return None, None
 
     key = random.randint(0, 999)
     
@@ -52,16 +56,16 @@ def query(prompt, is_negative=False, steps=35, 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 Mistral Nemo model
-    prompt = enhance_prompt(prompt)
+    # Enhance the prompt using the Mistral Nemo model if enabled
+    enhanced_prompt = enhance_prompt(prompt, enable_enhancement)
 
     # Add some extra flair to the prompt
-    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key} final prompt:\033[0m {prompt}')
+    final_prompt = f"{enhanced_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    print(f'\033[1mGeneration {key} final prompt:\033[0m {final_prompt}')
     
     # Prepare the payload for the API call, including width and height
     payload = {
-        "inputs": prompt,
+        "inputs": final_prompt,
         "is_negative": is_negative,
         "steps": steps,
         "cfg_scale": cfg_scale,
@@ -86,11 +90,11 @@ def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Ka
         # Convert the response content into an image
         image_bytes = response.content
         image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
-        return image
+        print(f'\033[1mGeneration {key} completed!\033[0m ({final_prompt})')
+        return image, enhanced_prompt
     except Exception as e:
         print(f"Error when trying to open the image: {e}")
-        return None
+        return None, None
 
 # CSS to style the app
 css = """
@@ -122,6 +126,11 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                 # Accordion for advanced settings
                 with gr.Row():
                     with gr.Accordion("Advanced Settings", open=False):
+                        enable_enhancement = gr.Checkbox(
+                            label="Enable Prompt Enhancement",
+                            value=True,
+                            description="Enable or disable automatic prompt enhancement"
+                        )
                         negative_prompt = gr.Textbox(
                             label="Negative Prompt", 
                             placeholder="What should not be in the image", 
@@ -150,11 +159,15 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         with gr.Row():
             image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
         
+        # Text output area to display the enhanced prompt
+        with gr.Row():
+            prompt_output = gr.Textbox(label="Enhanced Prompt", elem_id="prompt-output")
+        
         # Bind the button to the query function with all inputs
         text_button.click(
             query, 
-            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], 
-            outputs=image_output
+            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height, enable_enhancement], 
+            outputs=[image_output, prompt_output]
         )
 
 # Launch the Gradio app