Update app.py

app.py

@@ -89,7 +89,7 @@ def mistral_nemo_call(prompt, API_TOKEN, model="mistralai/Mistral-Nemo-Instruct-
         response += message.choices[0].delta.content
     return response
     
-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, use_mistral_nemo=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_style="generic", enhance_prompt_option=False, nemo_enhance_prompt_style="generic", use_mistral_nemo=False):
     # Determine which API URL to use
     api_url = API_URL_DEV if use_dev else API_URL
 
@@ -117,10 +117,10 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
 
     original_prompt = prompt
     if enhance_prompt_option:
-        prompt = enhance_prompt_v2(prompt, style="microscopic")
+        prompt = enhance_prompt_v2(prompt, style=enhance_prompt_style)
         print(f'\033[1mGeneration {key} enhanced prompt:\033[0m {prompt}')
     if use_mistral_nemo:
-        prompt = mistral_nemo_call(prompt,API_TOKEN=API_TOKEN,style="cartoon")
+        prompt = mistral_nemo_call(prompt, API_TOKEN=API_TOKEN, style=nemo_enhance_prompt_style)
         print(f'\033[1mGeneration {key} Mistral-Nemo prompt:\033[0m {prompt}')
         
     final_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
@@ -247,8 +247,11 @@ 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_style =  gr.Textbox(label="Enhance Prompt Style", placeholder="Enter style for the prompt enhancer here", elem_id="enhance-prompt-style")
                         enhance_prompt_option = gr.Checkbox(label="Enhance Prompt", value=False, elem_id="enhance-prompt-checkbox")
                         use_mistral_nemo = gr.Checkbox(label="Use Mistral Nemo", value=False, elem_id="use-mistral-checkbox")
+                        nemo_prompt_style =  gr.Textbox(label="Nemo Enhance Prompt Style", placeholder="Enter style for the prompt enhancer here", elem_id="nemo-enhance-prompt-style")
+                        
         with gr.Row():
             text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
         with gr.Row():
@@ -258,6 +261,6 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
             final_prompt_output = gr.Textbox(label="Final Prompt", elem_id="final-prompt-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, use_mistral_nemo], outputs=[image_output, seed_output, final_prompt_output])
+        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev, enhance_prompt_style,enhance_prompt_option, enhance_prompt_style, use_mistral_nemo], outputs=[image_output, seed_output, final_prompt_output])
 
 app.launch(show_api=True, share=False)