Add GPU detection

app.py

@@ -1,6 +1,7 @@
 import re, os
 from pathlib import Path
 import gradio as gr
+import torch
 
 from evodiff.pretrained import OA_DM_38M, D3PM_UNIFORM_38M, MSA_OA_DM_MAXSUB
 from evodiff.generate import generate_oaardm, generate_d3pm
@@ -8,16 +9,22 @@ from evodiff.generate_msa import generate_query_oadm_msa_simple
 from evodiff.conditional_generation import inpaint_simple, generate_scaffold
 
 
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+
 def make_uncond_seq(seq_len, model_type):
     if model_type == "EvoDiff-Seq-OADM 38M":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
-        tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, int(seq_len), batch_size=1, device='cpu')
+        tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, int(seq_len), batch_size=1, device=get_device())
     
     if model_type == "EvoDiff-D3PM-Uniform 38M":
         checkpoint = D3PM_UNIFORM_38M(return_all=True)
         model, collater, tokenizer, scheme, timestep, Q_bar, Q = checkpoint
-        tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, int(seq_len), batch_size=1, device='cpu')
+        tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, int(seq_len), batch_size=1, device=get_device())
 
     return generated_sequence
 
@@ -26,7 +33,7 @@ def make_cond_seq(seq_len, msa_file, n_sequences, model_type):
         checkpoint = MSA_OA_DM_MAXSUB()
         model, collater, tokenizer, scheme = checkpoint
         print(f"MSA File Path: {msa_file.name}")
-        tokeinzed_sample, generated_sequence  = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, int(n_sequences), seq_length=int(seq_len), device='cpu', selection_type='random')
+        tokeinzed_sample, generated_sequence  = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, int(n_sequences), seq_length=int(seq_len), device=get_device(), selection_type='random')
 
     return generated_sequence
     
@@ -34,7 +41,7 @@ def make_inpainted_idrs(sequence, start_idx, end_idx, model_type):
     if model_type == "EvoDiff-Seq":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
-        sample, entire_sequence, generated_idr = inpaint_simple(model, sequence, int(start_idx), int(end_idx), tokenizer=tokenizer, device='cpu')
+        sample, entire_sequence, generated_idr = inpaint_simple(model, sequence, int(start_idx), int(end_idx), tokenizer=tokenizer, device=get_device())
 
         generated_idr_output = {
             "original_sequence": sequence,
@@ -54,7 +61,7 @@ def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_ty
         # print("Folders in User Cache Directory:", os.listdir("/home/user/.cache"))
         start_idx = list(map(int, start_idx.strip('][').split(',')))
         end_idx = list(map(int, end_idx.strip('][').split(',')))
-        generated_sequence, new_start_idx, new_end_idx = generate_scaffold(model, pdb_code, start_idx, end_idx, scaffold_length, data_top_dir, tokenizer, device='cpu')
+        generated_sequence, new_start_idx, new_end_idx = generate_scaffold(model, pdb_code, start_idx, end_idx, scaffold_length, data_top_dir, tokenizer, device=get_device())
 
         generated_scaffold_output = {
             "generated_sequence": generated_sequence,
@@ -100,7 +107,7 @@ idr_app = gr.Interface(
                 ],
             outputs=["text"],
             title = "Inpainting IDRs",
-            description="Inpaining a new region inside a given sequence using the `EvoDiff-Seq` model."
+            description="Inpainting a new region inside a given sequence using the `EvoDiff-Seq` model."
             )
 
 scaffold_app = gr.Interface(
@@ -133,14 +140,14 @@ with gr.Blocks() as edapp:
         gr.TabbedInterface([
             usg_app,
             csg_app,
-            idr_app#,
-            # scaffold_app
+            idr_app,
+            scaffold_app
             ],
             [
                 "Unconditional sequence generation",
                 "Conditional generation",
-                "Inpainting IDRs"#,
-                # "Scaffolding functional motifs"
+                "Inpainting IDRs",
+                "Scaffolding functional motifs"
             ])