Fix: Resolve runtime errors and add MCP support

app.py

@@ -24,9 +24,17 @@ theme = gr.themes.Monochrome(
     primary_hue="gray",
 )
 
-## Function to get model from Hugging Face using token
 def get_model(model_name, token):
+    """
+    Logs into Hugging Face and loads a specified ESM3 model.
 
+    Args:
+        model_name (str): The name of the model to load from Hugging Face.
+        token (str): The Hugging Face authentication token.
+
+    Returns:
+        ESM3: The loaded ESM3 model moved to the CUDA device.
+    """
     if token:
         login(token=token)
 
@@ -38,14 +46,36 @@ def get_model(model_name, token):
     # model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
     return model
 
-## Function to get PDB data
 def get_pdb(pdb_id, chain_id):
+    """
+    Fetches a protein structure from the RCSB PDB database.
+
+    Args:
+        pdb_id (str): The 4-character PDB identifier.
+        chain_id (str): The specific chain identifier within the PDB file.
+
+    Returns:
+        ProteinChain: An object representing the specified protein chain.
+    """
     pdb = ProteinChain.from_rcsb(pdb_id, chain_id)
     # return [pdb.sequence, render_pdb(pdb.to_pdb_string())]
     return pdb
 
-## Function to generate rep for 3D structure
 def make_reps(res_start=None, res_end=None, main_color="whiteCarbon", highlight_color="redCarbon", main_style="cartoon", highlight_style="cartoon"):
+    """
+    Creates representation styles for 3D protein visualization with py3Dmol.
+
+    Args:
+        res_start (int, optional): The starting residue for highlighting. Defaults to None.
+        res_end (int, optional): The ending residue for highlighting. Defaults to None.
+        main_color (str, optional): The color for the main protein structure. Defaults to "whiteCarbon".
+        highlight_color (str, optional): The color for the highlighted region. Defaults to "redCarbon".
+        main_style (str, optional): The style for the main structure (e.g., 'cartoon'). Defaults to "cartoon".
+        highlight_style (str, optional): The style for the highlighted region. Defaults to "cartoon".
+
+    Returns:
+        list: A list of dictionary objects defining the py3Dmol styles.
+    """
     residue_range = f"{res_start}-{res_end}" if res_start != res_end else ""
 
     return [
@@ -72,8 +102,21 @@ def make_reps(res_start=None, res_end=None, main_color="whiteCarbon", highlight_
           "visible": True
       }]
 
-## Function to render 3D structure
 def render_pdb(pdb_id, chain_id, res_start, res_end, pdb_string=None):
+    """
+    Renders a PDB structure for display in a Gradio Molecule3D component.
+
+    Args:
+        pdb_id (str): The PDB ID, used for file naming.
+        chain_id (str): The chain ID, used for file naming.
+        res_start (int): The starting residue for highlighting.
+        res_end (int): The ending residue for highlighting.
+        pdb_string (str, optional): A string containing PDB data. If None,
+                                    it will be fetched from RCSB. Defaults to None.
+
+    Returns:
+        Molecule3D: A Gradio component object for 3D visualization.
+    """
     if pdb_string is None:
         pdb_string = get_pdb(pdb_id, chain_id).to_pdb_string()
     ## Write to temporary file and read back in to get the 3D structure
@@ -82,11 +125,25 @@ def render_pdb(pdb_id, chain_id, res_start, res_end, pdb_string=None):
 
     return Molecule3D(tmp_pdb.name, reps=make_reps(res_start=res_start, res_end=res_end))
 
-
-
-## Function for Scaffolding
 @spaces.GPU()
 def scaffold(model_name, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
+    """
+    Performs protein scaffolding by generating a new protein structure around a
+    functional motif from an existing protein.
+
+    Args:
+        model_name (str): The ESM3 model to use.
+        pdb_id (str): PDB ID of the source protein.
+        chain_id (str): Chain of the source protein.
+        motif_start (int): Starting residue of the motif.
+        motif_end (int): Ending residue of the motif.
+        prompt_length (int): Total length of the new protein sequence.
+        insert_size (int): Position where the motif will be inserted.
+
+    Returns:
+        list: A list of outputs for the Gradio interface, including sequences
+              and 3D structures.
+    """
     pdb = get_pdb(pdb_id, chain_id)
 
     ## Get motif sequence and atom37 positions
@@ -145,9 +202,23 @@ def scaffold(model_name, pdb_id, chain_id, motif_start, motif_end, prompt_length
         structure_new_highlight
     ]
 
-## Function for Secondary Structure Editing
 @spaces.GPU()
 def ss_edit(model_name, pdb_id, chain_id, region_start, region_end, shortened_region_length, shortening_ss8):
+    """
+    Edits the secondary structure of a protein, for example, by shortening a helix.
+
+    Args:
+        model_name (str): The ESM3 model to use.
+        pdb_id (str): PDB ID of the source protein.
+        chain_id (str): Chain of the source protein.
+        region_start (int): Starting residue of the region to edit.
+        region_end (int): Ending residue of the region to edit.
+        shortened_region_length (int): The new length of the edited region.
+        shortening_ss8 (str): The 8-state secondary structure string for the original protein.
+
+    Returns:
+        list: A list of outputs for the Gradio interface.
+    """
     pdb = get_pdb(pdb_id, chain_id)
     edit_region = np.arange(region_start, region_end)
 
@@ -192,9 +263,23 @@ def ss_edit(model_name, pdb_id, chain_id, region_start, region_end, shortened_re
         structure_new_highlight
         ]
 
-## Function for SASA Editing
 @spaces.GPU()
 def sasa_edit(model_name, pdb_id, chain_id, span_start, span_end, n_samples):
+    """
+    Edits a protein region to increase its solvent-accessible surface area (SASA).
+
+    Args:
+        model_name (str): The ESM3 model to use.
+        pdb_id (str): PDB ID of the source protein.
+        chain_id (str): Chain of the source protein.
+        span_start (int): Starting residue of the region to edit.
+        span_end (int): Ending residue of the region to edit.
+        n_samples (int): The number of new designs to generate.
+
+    Returns:
+        list: A list of outputs for the Gradio interface, including the best
+              generated structure.
+    """
     pdb = get_pdb(pdb_id, chain_id)
 
     structure_prompt = torch.full((len(pdb), 37, 3), torch.nan)
@@ -236,7 +321,7 @@ scaffold_app = gr.Interface(
     fn=scaffold,
     inputs=[
         gr.Dropdown(label="Model Name", choices=["esm3_sm_open_v1"], value="esm3_sm_open_v1", allow_custom_value=True),
-        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"), 
+        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"),
         gr.Textbox(value="1ITU", label = "PDB Code"),
         gr.Textbox(value="A", label = "Chain"),
         gr.Number(value=123, label="Motif Start"),
@@ -262,7 +347,7 @@ ss_app = gr.Interface(
     fn=ss_edit,
     inputs=[
         gr.Dropdown(label="Model Name", choices=["esm3_sm_open_v1"], value="esm3_sm_open_v1", allow_custom_value=True),
-        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"), 
+        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"),
         gr.Textbox(value = "7XBQ", label="PDB ID"),
         gr.Textbox(value = "A", label="Chain ID"),
         gr.Number(value=38, label="Edit Region Start"),
@@ -289,7 +374,7 @@ sasa_app = gr.Interface(
     fn=sasa_edit,
     inputs=[
         gr.Dropdown(label="Model Name", choices=["esm3_sm_open_v1"], value="esm3_sm_open_v1", allow_custom_value=True),
-        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"), 
+        # gr.Textbox(value = "hf_...", label="Hugging Face Token", type="password"),
         gr.Textbox(value = "1LBS", label="PDB ID"),
         gr.Textbox(value = "A", label="Chain ID"),
         gr.Number(value=105, label="Span Start"),
@@ -332,7 +417,7 @@ with gr.Blocks(theme=theme) as esm_app:
             - GitHub: https://github.com/evolutionaryscale/esm
             - HuggingFace Model: https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1
 
-            Spaces App By: [[Colby T. Ford](https://colbyford.com)] from [Tuple, The Cloud Genomics Company](https://tuple.xyz) 
+            Spaces App By: [[Colby T. Ford](https://colbyford.com)] from [Tuple, The Cloud Genomics Company](https://tuple.xyz)
 
             NOTE: You will need to agree to EvolutionaryScale's [license agreement](https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1) to use the model. This space uses a stored token for API access.
             """
@@ -352,4 +437,4 @@ with gr.Blocks(theme=theme) as esm_app:
             ])
 
 if __name__ == "__main__":
-    esm_app.launch()
+    esm_app.launch(mcp_server=True)