Fix: Resolve runtime errors by enabling GPU, refactoring auth and updating gradio

The Space was previously failing to run due to gradio issues and incorrect hardware requests. This commit resolves these runtime errors, making the application fully functional.

Authentication Fix: The user-facing token input, which was causing errors, has been removed. The application now correctly authenticates by reading the HF_TOKEN from the Space's secrets.
GPU Enablement: Added the @spaces.GPU() decorator to the main processing functions. This ensures the Space requests and utilizes the necessary GPU hardware, fixing a key execution failure.
Dependency Update: Corrected the Gradio SDK version in README.md to 5.38.2.
UI Update: The description text in the UI has been updated to reflect the new, functional authentication method.

README.md

@@ -4,7 +4,7 @@ emoji: 🧬
 colorFrom: gray
 colorTo: green
 sdk: gradio
-sdk_version: 4.37.1
+sdk_version: 5.38.2
 app_file: app.py
 pinned: true
 short_description: A frontier generative model for biology by EvolutionaryScale

app.py

@@ -1,11 +1,11 @@
-
+import spaces
 import gradio as gr
 import numpy as np
 import os, tempfile
 import torch
 import py3Dmol
 from huggingface_hub import login
-# import spaces
+
 
 from esm.utils.structure.protein_chain import ProteinChain
 from esm.models.esm3 import ESM3
@@ -16,20 +16,24 @@ from esm.sdk.api import (
 
 from gradio_molecule3d import Molecule3D
 
+# --- Retrieve the HF token from the Space's secrets ---
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 
 theme = gr.themes.Monochrome(
     primary_hue="gray",
 )
 
 ## Function to get model from Hugging Face using token
-# @spaces.GPU()
 def get_model(model_name, token):
-    login(token=token)
 
-    if torch.cuda.is_available():
-        model = ESM3.from_pretrained(model_name, device=torch.device("cuda"))
-    else:
-        model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
+    if token:
+        login(token=token)
+
+    # if torch.cuda.is_available():
+    model = ESM3.from_pretrained(model_name, device=torch.device("cuda"))
+    # else:
+    #     model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
 
     # model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
     return model
@@ -45,28 +49,28 @@ def make_reps(res_start=None, res_end=None, main_color="whiteCarbon", highlight_
     residue_range = f"{res_start}-{res_end}" if res_start != res_end else ""
 
     return [
-     {
-        "model": 0,
-        "chain": "",
-        "resname": "",
-        "style": main_style,
-        "color": main_color,
-        "residue_range": "",
-        "around": 0,
-        "byres": False,
-        "visible": True
-    },
-    {
-        "model": 0,
-        "chain": "",
-        "resname": "",
-        "style": highlight_style,
-        "color": highlight_color,
-        "residue_range": residue_range,
-        "around": 0,
-        "byres": False,
-        "visible": True
-    }]
+      {
+          "model": 0,
+          "chain": "",
+          "resname": "",
+          "style": main_style,
+          "color": main_color,
+          "residue_range": "",
+          "around": 0,
+          "byres": False,
+          "visible": True
+      },
+      {
+          "model": 0,
+          "chain": "",
+          "resname": "",
+          "style": highlight_style,
+          "color": highlight_color,
+          "residue_range": residue_range,
+          "around": 0,
+          "byres": False,
+          "visible": True
+      }]
 
 ## Function to render 3D structure
 def render_pdb(pdb_id, chain_id, res_start, res_end, pdb_string=None):
@@ -78,8 +82,11 @@ 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
-def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
+@spaces.GPU()
+def scaffold(model_name, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
     pdb = get_pdb(pdb_id, chain_id)
 
     ## Get motif sequence and atom37 positions
@@ -102,7 +109,7 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
                                                   num_steps=sequence_prompt.count("_") // 2,
                                                   temperature=0.5)
     ## Generate sequence
-    model = get_model(model_name, token)
+    model = get_model(model_name, HF_TOKEN)
     sequence_generation = model.generate(protein_prompt, sequence_generation_config)
     generated_sequence = sequence_generation.sequence
 
@@ -110,7 +117,7 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     structure_prediction_config = GenerationConfig(
         track="structure", # We want ESM3 to generate tokens for the structure track
         num_steps=len(sequence_generation) // 8,
-        temperature=0.7, 
+        temperature=0.7,
     )
     structure_prediction_prompt = ESMProtein(sequence=sequence_generation.sequence)
     structure_prediction = model.generate(structure_prediction_prompt, structure_prediction_config)
@@ -121,14 +128,14 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     # crmsd = structure_prediction_chain.rmsd(renal_dipep_chain, mobile_inds=motif_inds_in_generation, target_inds=motif_inds)
 
     structure_orig_highlight = render_pdb(pdb_id, chain_id, res_start=motif_start, res_end=motif_end)
-    structure_new_highlight = render_pdb(pdb_id, chain_id, res_start=insert_size, res_end=insert_size+len(motif_sequence), 
+    structure_new_highlight = render_pdb(pdb_id, chain_id, res_start=insert_size, res_end=insert_size+len(motif_sequence),
                                          pdb_string=structure_prediction_chain.to_pdb_string())
 
     return [
         pdb.sequence,
         motif_sequence,
         structure_orig_highlight,
-        # motif_atom37_positions,
+        # gr.Textbox(label="Motif Positions")
         sequence_prompt,
         # structure_prompt,
         # protein_prompt
@@ -139,7 +146,8 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     ]
 
 ## Function for Secondary Structure Editing
-def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, shortened_region_length, shortening_ss8):
+@spaces.GPU()
+def ss_edit(model_name, pdb_id, chain_id, region_start, region_end, shortened_region_length, shortening_ss8):
     pdb = get_pdb(pdb_id, chain_id)
     edit_region = np.arange(region_start, region_end)
 
@@ -148,21 +156,21 @@ def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, short
 
     ## Construct a secondary structure prompt that retains the secondary structure of the flanking regions, and shortens the lengths of helices in the helix-coil-helix region
     ss8_prompt = shortening_ss8[:edit_region[0]] + (((shortened_region_length - 3) // 2) * "H" + "C"*3 + ((shortened_region_length - 3) // 2) * "H") + shortening_ss8[edit_region[-1] + 1:]
-    
+
     ## Save original sequence and secondary structure
     original_sequence = pdb.sequence
     original_ss8 = shortening_ss8
     original_ss8_region = " "*edit_region[0] + shortening_ss8[edit_region[0]:edit_region[-1]+1]
-    
+
     proposed_ss8_region = " "*edit_region[0] + ss8_prompt[edit_region[0]:edit_region[0]+shortened_region_length]
 
     ## Create protein prompt
     protein_prompt = ESMProtein(sequence=sequence_prompt, secondary_structure=ss8_prompt)
 
     ## Generatre sequence
-    model = get_model(model_name, token)
+    model = get_model(model_name, HF_TOKEN)
     sequence_generation = model.generate(protein_prompt, GenerationConfig(track="sequence", num_steps=protein_prompt.sequence.count("_") // 2, temperature=0.5))
-    
+
     ## Generate structure
     structure_prediction = model.generate(ESMProtein(sequence=sequence_generation.sequence), GenerationConfig(track="structure", num_steps=len(protein_prompt) // 4, temperature=0))
     structure_prediction_chain = structure_prediction.to_protein_chain()
@@ -185,18 +193,19 @@ def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, short
         ]
 
 ## Function for SASA Editing
-def sasa_edit(model_name, token, pdb_id, chain_id, span_start, span_end, n_samples):
+@spaces.GPU()
+def sasa_edit(model_name, pdb_id, chain_id, span_start, span_end, n_samples):
     pdb = get_pdb(pdb_id, chain_id)
 
     structure_prompt = torch.full((len(pdb), 37, 3), torch.nan)
-    structure_prompt[span_start:span_end] = torch.tensor(pdb[span_start:span_end].atom37_positions, dtype=torch.float32)   
+    structure_prompt[span_start:span_end] = torch.tensor(pdb[span_start:span_end].atom37_positions, dtype=torch.float32)
 
     sasa_prompt = [None]*len(pdb)
     sasa_prompt[span_start:span_end] = [40.0]*(span_end - span_start)
 
     protein_prompt = ESMProtein(sequence="_"*len(pdb), coordinates=structure_prompt, sasa=sasa_prompt)
 
-    model = get_model(model_name, token)
+    model = get_model(model_name, HF_TOKEN)
 
     generated_proteins = []
     for i in range(n_samples):
@@ -227,7 +236,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 = "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", 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"),
@@ -253,7 +262,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 = "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", 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"),
@@ -280,12 +289,12 @@ 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 = "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", 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"),
         gr.Number(value=116, label="Span End"),
-        # gr.Textbox(value="CCSSCCCCSSCHHHHHHTEEETTBBTTBCSSEEEEECCTTCCHHHHHTTTHHHHHHHTTCEEEEECCTTTTCSCHHHHHHHHHHHHHHHHHHTTSCCEEEEEETHHHHHHHHHHHHCGGGGGTEEEEEEESCCTTCBGGGHHHHHTTCBCHHHHHTBTTCHHHHHHHHTTTTBCSSCEEEEECTTCSSSCCCCSSSTTSTTCCBTSEEEEHHHHHCTTCCCCSHHHHHBHHHHHHHHHHHHCTTSSCCGGGCCSTTCCCSBCTTSCHHHHHHHHSTHHHHHHHHHHSCCBSSCCCCCGGGGGGSTTCEETTEECCC", label="SS8 String")
+        # gr.Textbox(value="...", label="SS8 String")
         gr.Number(value=1, label="Number of Samples")
     ],
     outputs = [
@@ -325,7 +334,7 @@ with gr.Blocks(theme=theme) as esm_app:
 
             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. Then, create and paste your HuggingFace token in the appropriate field.
+            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.
             """
         )
     with gr.Row():
@@ -343,4 +352,4 @@ with gr.Blocks(theme=theme) as esm_app:
             ])
 
 if __name__ == "__main__":
-    esm_app.launch()
+    esm_app.launch()