add api call

app.py

@@ -1,16 +1,11 @@
 import gradio as gr
 import py3Dmol
 import io
-import torch
 import numpy as np
 import os
 import traceback
 import spaces
-# # Install ESM
-# os.system("pip install esm")
-
-from huggingface_hub import login
-from esm.models.esm3 import ESM3
+from esm.sdk import client
 from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
 from esm.utils.structure.protein_chain import ProteinChain
 from Bio.Data import PDBData
@@ -18,14 +13,16 @@ import biotite.structure as bs
 from biotite.structure.io import pdb
 from esm.utils import residue_constants as RC
 
-# Login to Hugging Face Hub 
-hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-if not hf_token:
-    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
-login(token=hf_token)
-
 # Initialize the model
-model: ESM3InferenceClient = ESM3.from_pretrained("esm3-open").to("cuda" if torch.cuda.is_available() else "cpu")
+token = os.environ.get("ESM_API_TOKEN")
+if not token:
+    raise ValueError("ESM_API_TOKEN environment variable is not set")
+
+model = client(
+    model="esm3-medium-2024-03",
+    url="https://forge.evolutionaryscale.ai",
+    token=token,
+)
 
 amino3to1 = {
     'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',
@@ -122,7 +119,7 @@ def get_protein(pdb_file) -> ESMProtein:
 def add_noise_to_coordinates(protein: ESMProtein, noise_level: float) -> ESMProtein:
     """Add Gaussian noise to the atom positions of the protein."""
     coordinates = protein.coordinates
-    noise = torch.randn_like(coordinates) * noise_level
+    noise = np.random.randn(*coordinates.shape) * noise_level
     noisy_coordinates = coordinates + noise
     return ESMProtein(sequence=protein.sequence, coordinates=noisy_coordinates)