Fix file upload handling in app.py

app.py

@@ -14,8 +14,24 @@ zero = torch.Tensor([0]).cuda() if torch.cuda.is_available() else torch.Tensor([
 print(f"Device: {zero.device}")
 
 @spaces.GPU
-def run_gaia(protein_content, ligand_content, use_cphmd=False, use_qmmm=False):
+def run_gaia(protein_file, ligand_file, use_cphmd=False, use_qmmm=False):
     try:
+        # Read the uploaded files
+        protein_content = protein_file
+        ligand_content = ligand_file
+        
+        # Check if files are file objects or bytes
+        if hasattr(protein_file, 'read'):
+            protein_content = protein_file.read()
+        if hasattr(ligand_file, 'read'):
+            ligand_content = ligand_file.read()
+        
+        # If bytes, decode to string
+        if isinstance(protein_content, bytes):
+            protein_content = protein_content.decode('utf-8')
+        if isinstance(ligand_content, bytes):
+            ligand_content = ligand_content.decode('utf-8')
+        
         with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False, mode='w') as p:
             p.write(protein_content)
             protein_path = p.name
@@ -66,12 +82,7 @@ def score_wrapper(protein_file, ligand_file, cphmd, qmmm):
     if protein_file is None or ligand_file is None:
         return "Please upload both protein and ligand files"
     
-    result = run_gaia(
-        protein_file.decode('utf-8') if isinstance(protein_file, bytes) else protein_file,
-        ligand_file.decode('utf-8') if isinstance(ligand_file, bytes) else ligand_file,
-        cphmd,
-        qmmm
-    )
+    result = run_gaia(protein_file, ligand_file, cphmd, qmmm)
     
     if "error" in result:
         return f"❌ Error: {result['error']}"
@@ -99,8 +110,8 @@ with gr.Blocks(title="GAIA Drug Discovery", theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column():
-            protein_input = gr.File(label="📁 Protein File (PDB/PDBQT)")
-            ligand_input = gr.File(label="📁 Ligand File (PDB/PDBQT)")
+            protein_input = gr.File(label="📁 Protein File (PDB/PDBQT)", file_types=[".pdb", ".pdbqt"])
+            ligand_input = gr.File(label="📁 Ligand File (PDB/PDBQT)", file_types=[".pdb", ".pdbqt"])
             
             with gr.Row():
                 cphmd_check = gr.Checkbox(label="🧬 Constant-pH MD", value=False)