Update app.py

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import os
 import sys
+import asyncio
 sys.path.append("/home/user/app")  # Ensures Hugging Face looks in the correct place
+
 from HuggingFaceHelper import HuggingFaceHelper
 from AICoreAGIX_with_TB import AICoreAGIX
 
@@ -9,11 +11,23 @@ from AICoreAGIX_with_TB import AICoreAGIX
 ai_core = AICoreAGIX()
 
 # Initialize Hugging Face training helper
-helper = HuggingFaceHelper(model_path="./merged_model")
+helper = HuggingFaceHelper(model_path="Raiff1982/codriao-finetuned")
 
-def diagnose_tb(image_file, audio_file):
+async def diagnose_tb_async(image_file, audio_file):
     user_id = 1  # Placeholder user ID
-    result = ai_core.run_tb_diagnostics(image_file.name, audio_file.name, user_id)
+
+    if image_file is None or audio_file is None:
+        return "Please upload both a TB saliva image and a cough audio file."
+
+    result = await ai_core.run_tb_diagnostics(image_file.name, audio_file.name, user_id)
+
+    # Optional file cleanup
+    try:
+        os.remove(image_file.name)
+        os.remove(audio_file.name)
+    except:
+        pass
+
     return (
         f"**TB Risk Level:** {result['tb_risk']}\n\n"
         f"**Image Result:** {result['image_analysis']['result']} "
@@ -25,7 +39,13 @@ def diagnose_tb(image_file, audio_file):
         f"**Shareable Link:** {result['shareable_link']}"
     )
 
+def diagnose_tb(image_file, audio_file):
+    return asyncio.run(diagnose_tb_async(image_file, audio_file))
+
 def upload_and_finetune(jsonl_file):
+    if jsonl_file is None:
+        return "Please upload a .jsonl file to fine-tune Codriao."
+
     save_path = f"./training_data/{jsonl_file.name}"
     os.makedirs("training_data", exist_ok=True)
 
@@ -36,7 +56,12 @@ def upload_and_finetune(jsonl_file):
     helper.dataset_path = save_path
     helper.fine_tune(output_dir="./codette_finetuned")
 
-    return f"✅ Fine-tuning complete! Model updated and stored."
+    try:
+        os.remove(save_path)
+    except:
+        pass
+
+    return "✅ Fine-tuning complete! Model updated and stored."
 
 def get_latest_model():
     return "Download the latest fine-tuned Codriao model here: https://huggingface.co/Raiff1982/codriao-finetuned"