Update app.py

app.py

@@ -16,15 +16,35 @@ API_URL_DEV = "https://api-inference.huggingface.co/models/black-forest-labs/FLU
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 timeout = 100
 
-def enhance_prompt(prompt, style="photorealistic"):
+def enhance_prompt(prompt, model="mistralai/Mistral-7B-Instruct-v0.1", style="photo-realistic"):
+    
     client = Client("K00B404/Mistral-Nemo-custom")
+    
+    system_prompt=f"""
+    You are a image generation prompt enhancer specialized in the {style} style. 
+    You must respond only with the enhanced version of the users input prompt
+    Remember, image generation models can be stimulated by refering to camera 'effect' in the prompt like :4k ,award winning, super details, 35mm lens, hd
+    """
+    user_message=f"###input image generation prompt### {prompt}"
+   
     result = client.predict(
-        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"
+    		system_prompt=system_prompt,
+    		user_message=user_message,
+    		max_tokens=256,
+    		model_id=model,# "mistralai/Mistral-Nemo-Instruct-2407",
+    		api_name="/predict"
     )
     return result
     
+    # The output value that appears in the "Response" Textbox component.
+    """result = client.predict(
+        system_prompt=system_prompt,#"You are a image generation prompt enhancer and must respond only with the enhanced version of the users input prompt",
+        user_message=user_message,
+        max_tokens=500,
+        api_name="/predict"
+    )
+    return result
+    """
 def mistral_nemo_call(prompt, API_TOKEN, model="mistralai/Mistral-Nemo-Instruct-2407", style="photo-realistic"):
     
     client = InferenceClient(api_key=API_TOKEN)
@@ -74,7 +94,7 @@ 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(prompt)
+        prompt = enhance_prompt(prompt, style="cartoon")
         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")