Update app.py

app.py

@@ -14,24 +14,24 @@ pixtral_client = InferenceClient("mistralai/Pixtral-12B-2409")
 
 def enhance_prompt(prompt: str) -> str:
     # Use the Paligemma models for prompt enhancement
-    prompt_224 = paligemma224_client.query(inputs={"inputs": prompt})["generated_text"]
-    prompt_448 = paligemma448_client.query(inputs={"inputs": prompt})["generated_text"]
-    prompt_896 = paligemma896_client.query(inputs={"inputs": prompt})["generated_text"]
+    prompt_224 = paligemma224_client(prompt)["generated_text"]
+    prompt_448 = paligemma448_client(prompt)["generated_text"]
+    prompt_896 = paligemma896_client(prompt)["generated_text"]
     
     # Combine all enhanced prompts into a single one
     enhanced_prompt = f"Enhanced (224): {prompt_224}\nEnhanced (448): {prompt_448}\nEnhanced (896): {prompt_896}"
     
     # Ultra-enhance the prompt using Paligemma 28b
-    ultra_enhanced_prompt = paligemma28b_client.query(inputs={"inputs": enhanced_prompt})["generated_text"]
+    ultra_enhanced_prompt = paligemma28b_client(enhanced_prompt)["generated_text"]
     
     return ultra_enhanced_prompt
 
 
 def generate_answer(enhanced_prompt: str) -> str:
     # Generate answers using the three models: llama, deepseek, and omniparser
-    llama_answer = llama_client.query(inputs={"inputs": enhanced_prompt})["generated_text"]
-    deepseek_answer = deepseek_client.query(inputs={"inputs": enhanced_prompt})["generated_text"]
-    omniparser_answer = omniparser_client.query(inputs={"inputs": enhanced_prompt})["generated_text"]
+    llama_answer = llama_client(enhanced_prompt)["generated_text"]
+    deepseek_answer = deepseek_client(enhanced_prompt)["generated_text"]
+    omniparser_answer = omniparser_client(enhanced_prompt)["generated_text"]
     
     # Combine answers from all models
     combined_answer = f"Llama: {llama_answer}\nDeepseek: {deepseek_answer}\nOmniparser: {omniparser_answer}"
@@ -41,9 +41,10 @@ def generate_answer(enhanced_prompt: str) -> str:
 
 def enhance_answer(answer: str) -> str:
     # Enhance the final answer using Pixtral model
-    enhanced_answer = pixtral_client.query(inputs={"inputs": answer})["generated_text"]
+    enhanced_answer = pixtral_client(answer)["generated_text"]
     return enhanced_answer
 
+
 def process(message: str) -> str:
     # Step 1: Enhance the prompt using the Paligemma models
     enhanced_prompt = enhance_prompt(message)