Update app.py

app.py

@@ -103,6 +103,24 @@ def mistral_nemo_call(prompt, API_TOKEN, model="mistralai/Mistral-Nemo-Instruct-
     return response
 
 
+
+def chat_with_persona(message, history, system_message, max_tokens, temperature, top_p):
+    """Function to interact with the chatbot API using the generated persona"""
+    try:
+        # Call the API with the current message and system prompt (persona)
+        response = chatbot_client.predict(
+            message=message,
+            system_message=system_message,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            api_name="/chat"
+        )
+        return response
+    except Exception as e:
+        return f"Error communicating with the chatbot API: {str(e)}"
+        
+
 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, 
@@ -137,10 +155,25 @@ 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=enhance_prompt_style)
+        style=enhance_prompt_style
+        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
+    """
+        prompt = chat_with_persona(message=prompt, history=[], system_message=system_prompt, max_tokens=256, temperature=0.1, top_p=0.97)
+        #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=nemo_enhance_prompt_style)
+        style=nemo_enhance_prompt_style
+        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
+    """
+        prompt = chat_with_persona(message=prompt, history=[], system_message=system_prompt, max_tokens=256, temperature=0.1, top_p=0.97)
+        #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."