Update app.py

app.py

@@ -15,7 +15,7 @@ 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):
+def enhance_prompt(prompt, style="photorealistic"):
     client = Client("K00B404/Mistral-Nemo-custom")
     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",
@@ -23,7 +23,26 @@ def enhance_prompt(prompt):
         api_name="/predict"
     )
     return result
-
+    
+def mistral_nemo_call(prompt, model="mistralai/Mistral-Nemo-Instruct-2407"):
+    
+    client = InferenceClient(api_key=API_TOKEN)
+    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
+    """,
+    
+    response = ""
+    for message in client.chat_completion(
+        model="mistralai/Mistral-Nemo-Instruct-2407",
+        messages=[{"role": "system", "content": system_prompt},{"role": "user", "content": prompt}],
+        max_tokens=500,
+        stream=True,
+    ):
+        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):
     # Determine which API URL to use
     api_url = API_URL_DEV if use_dev else API_URL
@@ -34,12 +53,13 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     if is_api_call:
         # Use the environment variable for the API key in GUI mode
         API_TOKEN = os.getenv("HF_READ_TOKEN")
-        headers = {"Authorization": f"Bearer {API_TOKEN}"}
     else:
         # Validate the API key if it's an API call
         if huggingface_api_key == "":
             raise gr.Error("API key is required for API calls.")
-        headers = {"Authorization": f"Bearer {huggingface_api_key}"}
+        API_TOKEN = huggingface_api_key
+    
+    headers = {"Authorization": f"Bearer {API_TOKEN}"} 
 
     if prompt == "" or prompt is None:
         return None, None, None
@@ -53,7 +73,10 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     if enhance_prompt_option:
         prompt = enhance_prompt(prompt)
         print(f'\033[1mGeneration {key} enhanced prompt:\033[0m {prompt}')
-
+    elif use_mistral_nemo:
+        prompt = mistral_nemo_call(prompt)
+        print(f'\033[1mGeneration {key} Mistral-Nemo prompt:\033[0m {prompt}')
+        
     final_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'\033[1mGeneration {key}:\033[0m {final_prompt}')