Update app.py

app.py

@@ -13,7 +13,7 @@ TOKEN = os.getenv("HF_TOKEN")
 login(TOKEN)
 
 # This will download the model weights and instantiate the model on your machine.
-model: ESM3InferenceClient = ESM3.from_pretrained("esm3_sm_open_v1").to("cpu")
+model: ESM3InferenceClient = ESM3.from_pretrained("esm3_sm_open_v1").to("cuda")
 
 def read_mol(molpath):
     with open(molpath, "r") as fp:
@@ -75,30 +75,38 @@ def molecule(input_pdb):
     allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
     allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
 
-def prediction(prompt, temperature):
+def prediction(prompt, temperature, do_structure, enable_roundtrip):
     protein = ESMProtein(sequence=prompt)
     protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8, temperature=temperature))
-    seq = protein.sequence
-    protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
-    protein.to_pdb("./generation.pdb")   
-    protein.sequence = None
-    protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8))
-    protein.coordinates = None
-    protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
-    protein.to_pdb("./round_tripped.pdb")
-    html1 = molecule("./round_tripped.pdb")
-    html = molecule("./generation.pdb")
-    return seq, protein.sequence, html, html1, "./generation.pdb", "./round_tripped.pdb"
+    if do_structure == "Yes":
+        protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
+        protein.to_pdb("./generation.pdb")   
+        html = molecule("./generation.pdb")
+        if enable_roundtrip == "Yes":
+            seq = protein.sequence
+            protein.sequence = None
+            protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8))
+            protein.coordinates = None
+            protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
+            protein.to_pdb("./round_tripped.pdb")
+            html1 = molecule("./round_tripped.pdb")
+            return seq, protein.sequence, html, html1, "./generation.pdb", "./round_tripped.pdb"
+        else:
+            html1 = "<h3>Inverse folding and re-generation not enabled</h3>"
+            f = open("./round_tripped.pdb", "w")
+            f.write("\n")
+            f.close()
+            return protein.sequence, "Inverse folding and re-generation not enabled", html, html1, "./generation.pdb", "./round_tripped.pdb"
+    else:
+        f = open("./empty.pdb", "w")
+        f.write("\n")
+        f.close()
+        return protein.sequence, "Inverse folding and re-generation not enabled", "<h3>Structure reconstruction not enabled</h3>", "<h3>Inverse folding and re-generation not enabled</h3>", "./empty.pdb", "./empty.pdb"
+    
 
-def return_stuff(prompt, temperature):
-    seq, sequence, html, html1, pdb, pdb1 = prediction(prompt, temperature)
-    return seq, sequence, html, html1, pdb, pdb1
-
-demo = gr.Interface(fn = return_stuff, inputs = [gr.Textbox(label="Masked protein sequence", info="Use '_' as masking character", value="___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"), gr.Slider(0,1,label="Temperature")], outputs = [gr.Textbox(label="Originally predicted sequence"),gr.Textbox(label="Inverse folding predicted sequence"),gr.HTML(label="Predicted 3D structure"),gr.HTML(label="Inverse-folding predicted 3D structure"), Molecule3D(label="Predicted molecular structure"), Molecule3D(label="Inverse-folding predicted molecular structure")], title="""<h1 align='center'>Proteins with ESM</h1>
+demo = gr.Interface(fn = prediction, inputs = [gr.Textbox(label="Masked protein sequence", info="Use '_' as masking character", value="___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"), gr.Slider(0,1,label="Temperature"), gr.Radio(["Yes", "No"], label="Reconstruct structure", info="Choose wheter to reconstruct structure or not"), gr.Radio(["Yes", "No"], label="Allow inverse-folding", info="Choose wether to allow double check of prediction with inverse folding (ONLY when 'Reconstruct structure' is set to 'Yes')")], outputs = [gr.Textbox(label="Originally predicted sequence"),gr.Textbox(label="Inverse folding predicted sequence"),gr.HTML(label="Predicted 3D structure"),gr.HTML(label="Inverse-folding predicted 3D structure"), Molecule3D(label="Predicted molecular structure"), Molecule3D(label="Inverse-folding predicted molecular structure")], title="""<h1 align='center'>Proteins with ESM</h1>
     <h2 align='center'>Predict the whole sequence and 3D structure of masked protein sequences!</h2>
     <h3 align='center'>Support this space with a ⭐ on <a href='https://github.com/AstraBert/proteins-w-esm'>GitHub</a></h3>
-    <h3 align='center'>Support Evolutionary Scale's ESM with a ⭐ on <a href='https://github.com/evolutionaryscale/esm'>GitHub</a></h3>""", examples = [["___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________", 0.7], ["__________________________________________________________AGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVRE___________________________", 0.2], ["__________KTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLH________", 0.5]], cache_examples=False)
-
+    <h3 align='center'>Support Evolutionary Scale's ESM with a ⭐ on <a href='https://github.com/evolutionaryscale/esm'>GitHub</a></h3>""", examples = [["___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________", 0.7, "No", "No"], ["__________________________________________________________AGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVRE___________________________", 0.2, "Yes", "No"], ["__________KTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLH________", 0.5, "Yes", "Yes"]], cache_examples=False)
 
-if __name__ == '__main__':
-    demo.launch()
+demo.launch()