Update app.py

app.py

@@ -25,7 +25,7 @@ def enhance_prompt(prompt, style="photorealistic"):
     )
     return result
     
-def mistral_nemo_call(prompt, model="mistralai/Mistral-Nemo-Instruct-2407"):
+def mistral_nemo_call(prompt, API_TOKEN, model="mistralai/Mistral-Nemo-Instruct-2407"):
     
     client = InferenceClient(api_key=API_TOKEN)
     system_prompt=f"""
@@ -36,7 +36,7 @@ def mistral_nemo_call(prompt, model="mistralai/Mistral-Nemo-Instruct-2407"):
     
     response = ""
     for message in client.chat_completion(
-        model="mistralai/Mistral-Nemo-Instruct-2407",
+        model=model,
         messages=[{"role": "system", "content": system_prompt},{"role": "user", "content": prompt}],
         max_tokens=500,
         stream=True,
@@ -75,7 +75,7 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
         prompt = enhance_prompt(prompt)
         print(f'\033[1mGeneration {key} enhanced prompt:\033[0m {prompt}')
     if use_mistral_nemo:
-        prompt = mistral_nemo_call(prompt)
+        prompt = mistral_nemo_call(prompt,API_TOKEN=API_TOKEN)
         print(f'\033[1mGeneration {key} Mistral-Nemo prompt:\033[0m {prompt}')
         
     final_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."