Upload app.py

app.py

@@ -1,13 +1,16 @@
 import gradio as gr # type: ignore
 import os
-from gradio_molecule3d import Molecule3D
-from Bio.PDB import PDBParser
+from gradio_molecule3d import Molecule3D #type: ignore
+from Bio.PDB import PDBParser #type: ignore
+
+import time
 
 # Import your custom modules from the /scripts folder
 from scripts.download import download_and_clean_pdb
 from scripts.generator import run_broteinshake_generator
 from scripts.refine import polish_design, process_results
 from scripts.visualize import create_design_plot
+from scripts.foldprotein import fold_protein_sequence
 
 # --- HELPER FUNCTIONS ---
 
@@ -42,22 +45,22 @@ def load_pdb_and_extract_chains(pdb_id):
         # Single-chain proteins are supported (will use different ProteinMPNN command)
         # For single-chain, the only chain will be automatically selected for redesign
         if len(chains) == 1:
-            status_msg = f"✅ Loaded {pdb_id.upper()}: Single-chain protein - will redesign chain {chains[0]}"
+            status_msg = f"Loaded {pdb_id.upper()}: Single-chain protein - will redesign chain {chains[0]}"
             # Auto-select the chain for single-chain proteins
             return gr.update(choices=chains, value=chains), status_msg, gr.update(interactive=True), chains
         
-        status_msg = f"✅ Loaded {pdb_id.upper()}: Found {len(chains)} chain(s) - {', '.join(chains)}"
+        status_msg = f"Loaded {pdb_id.upper()}: Found {len(chains)} chain(s) - {', '.join(chains)}"
         # Initially disable button - user must select at least one chain
         return gr.update(choices=chains, value=chains), status_msg, gr.update(interactive=False), chains
     except Exception as e:
-        error_msg = f"❌ Error loading {pdb_id.upper()}: {str(e)}"
+        error_msg = f"Error loading {pdb_id.upper()}: {str(e)}"
         print(error_msg)
         return gr.update(choices=[], value=[]), error_msg, gr.update(interactive=False), []
 
 def validate_chain_selection(selected_chains, available_chains_state):
     """Validate that at least one chain is selected and at least one remains fixed (for multi-chain)."""
     if not selected_chains or len(selected_chains) == 0:
-        warning = "⚠️ Please select at least one chain to redesign"
+        warning = "Please select at least one chain to redesign"
         return gr.update(interactive=False), warning, available_chains_state
     
     # Get available chains from state
@@ -65,15 +68,15 @@ def validate_chain_selection(selected_chains, available_chains_state):
     
     # For single-chain proteins, allow selecting the only chain
     if len(available_chains) == 1:
-        warning = f"✅ Single-chain protein: Will redesign chain {available_chains[0]}"
+        warning = f"Single-chain protein: Will redesign chain {available_chains[0]}"
         return gr.update(interactive=True), warning, available_chains_state
     
     # For multi-chain: Check if all chains are selected (would leave no fixed chains)
     if available_chains and len(selected_chains) >= len(available_chains):
-        warning = f"⚠️ Cannot select all chains - at least one chain must remain fixed. Selected: {', '.join(selected_chains)}"
+        warning = f"Cannot select all chains - at least one chain must remain fixed. Selected: {', '.join(selected_chains)}"
         return gr.update(interactive=False), warning, available_chains_state
     
-    warning = f"✅ {len(selected_chains)} chain(s) selected for redesign: {', '.join(selected_chains)}"
+    warning = f"{len(selected_chains)} chain(s) selected for redesign: {', '.join(selected_chains)}"
     return gr.update(interactive=True), warning, available_chains_state
 
 def get_all_sequences(fasta_file: str) -> str:
@@ -160,9 +163,9 @@ def run_part1(pdb_id, fixed_chains, variable_chains, temperature=0.1, selected_c
                 raise ValueError(f"Cannot redesign all chains - at least one chain must remain fixed. Selected: {', '.join(selected_chains)}, Available: {', '.join(all_chains)}")
             
             if is_single_chain:
-                print(f"📋 Single-chain mode: Redesigning chain {variable_chains}")
+                print(f"Single-chain mode: Redesigning chain {variable_chains}")
             else:
-                print(f"📋 Using chain selector: Fixed={fixed_chains}, Variable={variable_chains}")
+                print(f"Using chain selector: Fixed={fixed_chains}, Variable={variable_chains}")
         else:
             # If no chains selected, use text inputs (default behavior)
             # For single-chain, if variable_chains is empty, use the only chain
@@ -174,12 +177,12 @@ def run_part1(pdb_id, fixed_chains, variable_chains, temperature=0.1, selected_c
                 all_selected = set(fixed_chains + variable_chains)
                 if len(all_selected) >= len(all_chains):
                     raise ValueError(f"Cannot redesign all chains - at least one chain must remain fixed.")
-            print(f"📋 Using text inputs: Fixed={fixed_chains}, Variable={variable_chains}")
+            print(f"Using text inputs: Fixed={fixed_chains}, Variable={variable_chains}")
         
         # Step 2: Generate Optimized Sequences
         # This creates the .fa files you need for the ESM Atlas
-        print(f"🌡️  Temperature: {temperature}")
-        print(f"🔧 Parameters: Fixed chains={fixed_chains}, Variable chains={variable_chains}, Temp={temperature}")
+        print(f"Temperature: {temperature}")
+        print(f"Parameters: Fixed chains={fixed_chains}, Variable chains={variable_chains}, Temp={temperature}")
         run_broteinshake_generator(pdb_path, fixed_chains, variable_chains, num_seqs=20, temp=temperature)
         
         # Get all sequences and the best one
@@ -199,58 +202,72 @@ def run_part1(pdb_id, fixed_chains, variable_chains, temperature=0.1, selected_c
         
         # Format status with best sequence
         status_message = (
-            f"✅ Design Complete! {num_designs} designs generated.\n\n"
-            f"⭐ Lead Candidate (Best Score: {best_score:.4f}):\n"
+            f"Design Complete! {num_designs} designs generated.\n\n"
+            f"Lead Candidate (Best Score: {best_score:.4f}):\n"
             f"{best_sequence}\n\n"
-            f"📋 Copy the lead sequence above and fold it at [esmatlas.com](https://esmatlas.com/resources?action=fold)."
+            
         )
             
         return all_sequences, evolution_plot, status_message
     except Exception as e:
-        return "", None, f"❌ Error in Part 1: {str(e)}"
+        return "", None, f"Error in Part 1: {str(e)}"
 
-def run_part2(pdb_id, uploaded_esm_file):
-    """Aligns the ESM-folded PDB to the target structure."""
+def run_part2(pdb_id, sequence):
+    """
+    1. Folds the input sequence using ESM Atlas (API).
+    2. Aligns the folded structure to the target PDB (polish_design).
+    """
     try:
-        if uploaded_esm_file is None:
-            return None, "⚠️ Please upload the PDB from ESM Atlas first."
+        # --- 1. Validate Inputs ---
+        if not sequence or not sequence.strip():
+            return None, "Error: Please enter a protein sequence."
         
-        # Validate inputs
         if not pdb_id or not pdb_id.strip():
-            return None, "❌ Error: PDB ID is required."
+            return None, "Error: PDB ID is required."
+        
+        print(f"Starting Pipeline for {pdb_id}...")
+        print(f"   - Sequence Length: {len(sequence)} residues")
+
+        # --- 2. Fold Sequence (Automated) ---
+        # Calls the script in scripts/foldprotein.py
+        pdb_content = fold_protein_sequence(sequence)
         
-        # Get file path
-        file_path = uploaded_esm_file.name if hasattr(uploaded_esm_file, 'name') else uploaded_esm_file
-        if not file_path or not os.path.exists(file_path):
-            return None, f"❌ Error: Uploaded file not found: {file_path}"
+        if not pdb_content:
+            return None, "Error: Protein folding failed. The API might be down or the sequence is invalid."
+
+        # Save the raw folded structure to a temp file
+        raw_fold_path = f"temp_fold_{pdb_id}_{int(time.time())}.pdb"
+        with open(raw_fold_path, "w") as f:
+            f.write(pdb_content)
+        
+        if not os.path.exists(raw_fold_path):
+             return None, f"Error: Could not save folded PDB to {raw_fold_path}"
+
+        # --- 3. Align & Polish (Existing Logic) ---
+        print(f"Aligning folded structure to target {pdb_id}...")
         
-        # Call the new lightweight Biopython-only script
-        print(f"🔍 Starting alignment for PDB ID: {pdb_id}, File: {file_path}")
-        final_pdb_path, global_rmsd, core_rmsd, high_conf_rmsd = polish_design(pdb_id, file_path)
+        # Pass the generated file path to the existing polish_design function
+        final_pdb_path, global_rmsd, core_rmsd, high_conf_rmsd = polish_design(pdb_id, raw_fold_path)
         
-        # Validate output
+        # --- 4. Validate Alignment Output ---
         if not final_pdb_path or not os.path.exists(final_pdb_path):
-            return None, f"❌ Error: Alignment failed - output file not created: {final_pdb_path}"
+            return None, f"Error: Alignment failed - output file not created: {final_pdb_path}"
         
         if high_conf_rmsd is None:
-            return None, "❌ Error: Alignment failed - RMSD calculation returned None"
+            return None, "Error: Alignment failed - RMSD calculation returned None"
         
-        print(f"✅ Alignment successful: Global RMSD = {global_rmsd:.3f} Å, Core RMSD = {core_rmsd:.3f} Å, High-Conf RMSD = {high_conf_rmsd:.3f} Å")
+        print(f"Success: Global RMSD={global_rmsd:.3f}A | Core RMSD={core_rmsd:.3f}A")
         
-        # Generate detailed validation report
+        # --- 5. Generate Report ---
         report = process_results(pdb_id, final_pdb_path, global_rmsd, high_conf_rmsd)
         
+        # Clean up the raw unaligned fold to save space
+        os.remove(raw_fold_path)
+        
         return final_pdb_path, report
-    except FileNotFoundError as e:
-        error_msg = f"❌ File Error: {str(e)}"
-        print(error_msg)
-        return None, error_msg
-    except ValueError as e:
-        error_msg = f"❌ Validation Error: {str(e)}"
-        print(error_msg)
-        return None, error_msg
+
     except Exception as e:
-        error_msg = f"❌ Unexpected Error: {str(e)}\n\nTraceback:\n{type(e).__name__}"
+        error_msg = f"Unexpected Error: {str(e)}"
         print(error_msg)
         import traceback
         traceback.print_exc()
@@ -297,7 +314,7 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
     # Header
     gr.HTML("""
         <div id='biohub-header'>
-            <h1 style='color: white; margin: 0;'>🧪 BroteinShake</h1>
+            <h1 style='color: white; margin: 0;'>BroteinShake</h1>
         </div>
     """)
     
@@ -307,11 +324,11 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
             gr.Markdown("Enter a PDB ID to 'repaint' its binder interface using ProteinMPNN.")
             
             pdb_input = gr.Textbox(label="Target PDB ID", placeholder="e.g., 3kas", value="")
-            load_pdb_btn = gr.Button("📥 Load PDB", variant="secondary")
-            pdb_status = gr.Markdown("💡 Enter a PDB ID and click 'Load PDB' to begin")
+            load_pdb_btn = gr.Button("Load PDB", variant="secondary")
+            pdb_status = gr.Markdown("Enter a PDB ID and click 'Load PDB' to begin")
             
             with gr.Column():
-                gr.Markdown("### ⚙️ Design Parameters")
+                gr.Markdown("### Design Parameters")
                 
                 # Temperature (T) is the most critical knob for sequence recovery
                 sampling_temp = gr.Slider(
@@ -327,7 +344,7 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
                     info="Identify which chains ProteinMPNN should modify (will populate after loading PDB)"
                 )
                 
-                chain_warning = gr.Markdown("💡 Select at least one chain to enable generation", visible=True)
+                chain_warning = gr.Markdown("Select at least one chain to enable generation", visible=True)
                 
                 # Hidden state to track if we've successfully parsed the PDB
                 pdb_state = gr.State()
@@ -338,7 +355,7 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
                     v_chains = gr.Textbox(label="Variable Chains (Key)", value="B")
             
             # Generate button (initially disabled)
-            gen_btn = gr.Button("🚀 Generate Optimized Sequences", variant="primary", interactive=False)
+            gen_btn = gr.Button("Generate Optimized Sequences", variant="primary", interactive=False)
             
             # Load PDB and extract chains when button is clicked
             load_pdb_btn.click(
@@ -373,7 +390,7 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
 
         # TAB 2: STRUCTURAL VALIDATION
         with gr.Tab("2. Structural Validation"):
-            gr.Markdown("### 🧬 Final Structure Preview")
+            gr.Markdown("### Final Structure Preview")
             
             # Updated REPS for local dev visibility
             REPS = [
@@ -382,48 +399,54 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
                     "style": "cartoon",
                     "color": "spectrum",
                     "opacity": 1.0
-                    # Removing "chain": "A" is the key fix here!
                 }
             ]
             
-            # This component ACTUALLY supports .pdb natively
+            # 3D Viewer component
             protein_view = Molecule3D(label="3D Structure Viewer (Refined Shuttle)", reps=REPS, elem_id="molecule-viewer")
             
             with gr.Row():
-                esm_upload = gr.File(label="Upload ESM-Folded PDB", file_types=[".pdb"])
+                # REPLACEMENT: Text Area instead of File Upload
+                sequence_input = gr.Textbox(
+                    label="Paste Protein Sequence", 
+                    placeholder="Paste the sequence generated in Tab 1 here (e.g., MKTII...)",
+                    lines=4,
+                    max_lines=8
+                )
                 refined_download = gr.File(label="Download Aligned Lead (.pdb)")
             
-            validate_btn = gr.Button("✨ Run Structural Alignment", variant="primary")
+            validate_btn = gr.Button("Run Structural Alignment", variant="primary")
             status2 = gr.Textbox(label="Validation Report", interactive=False, lines=5)
 
-            # Update viewer - REPS are set on component creation, just return the file path
-            def run_validation_with_view(pdb_id, file):
+            # Wrapper function now accepts 'sequence' instead of 'file'
+            def run_validation_with_view(pdb_id, sequence):
                 try:
-                    # This is your existing alignment/refine logic
-                    final_pdb, report = run_part2(pdb_id, file)
+                    # Call the updated run_part2 logic (API Fold -> Align)
+                    final_pdb, report = run_part2(pdb_id, sequence)
                     
                     if final_pdb is not None and os.path.exists(final_pdb):
-                        # Verify we're using the refined shuttle (Refined_Shuttle.pdb)
+                        # Verify we're using the refined shuttle
                         if "Refined_Shuttle.pdb" in final_pdb or os.path.basename(final_pdb) == "Refined_Shuttle.pdb":
-                            print(f"🎯 Visualizing refined shuttle: {final_pdb}")
+                            print(f"Visualizing refined shuttle: {final_pdb}")
                         else:
-                            print(f"⚠️ Warning: Expected Refined_Shuttle.pdb but got: {final_pdb}")
+                            print(f"Warning: Expected Refined_Shuttle.pdb but got: {final_pdb}")
                     
-                    # Molecule3D expects just the file path string (REPS are already set on component)
+                    # Return path for download, report text, and path for 3D viewer
                     return final_pdb, report, final_pdb
+                    
                 except Exception as e:
-                    error_msg = f"❌ Error generating 3D view: {str(e)}"
+                    error_msg = f"Error generating 3D view: {str(e)}"
                     print(error_msg)
                     import traceback
                     traceback.print_exc()
                     return None, error_msg, None
 
+            # Updated inputs to include sequence_input
             validate_btn.click(
                 run_validation_with_view, 
-                inputs=[pdb_input, esm_upload], 
+                inputs=[pdb_input, sequence_input], 
                 outputs=[refined_download, status2, protein_view]
             )
-
 # Launch the app
 if __name__ == "__main__":
     # Docker deployment for HuggingFace Spaces