ZeroGPU test 3

app.py

@@ -9,24 +9,19 @@ from evodiff.generate import generate_oaardm, generate_d3pm
 from evodiff.generate_msa import generate_query_oadm_msa_simple
 from evodiff.conditional_generation import inpaint_simple, generate_scaffold
 
-@spaces.GPU
-def get_device():
-    if torch.cuda.is_available():
-        return "cuda"
-    else:
-        return "cpu"
-
-@spaces.GPU
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+@spaces.GPU()
 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=get_device())
+        tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, int(seq_len), batch_size=1, device=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=get_device())
+        tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, int(seq_len), batch_size=1, device=device)
 
     return generated_sequence
 
@@ -35,7 +30,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=get_device(), 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=device, selection_type='random')
 
     return generated_sequence
     
@@ -43,7 +38,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=get_device())
+        sample, entire_sequence, generated_idr = inpaint_simple(model, sequence, int(start_idx), int(end_idx), tokenizer=tokenizer, device=device)
 
         generated_idr_output = {
             "original_sequence": sequence,
@@ -63,7 +58,7 @@ def make_inpainted_idrs(sequence, start_idx, end_idx, model_type):
 #         # 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=get_device())
+#         generated_sequence, new_start_idx, new_end_idx = generate_scaffold(model, pdb_code, start_idx, end_idx, scaffold_length, data_top_dir, tokenizer, device=device)
 
 #         generated_scaffold_output = {
 #             "generated_sequence": generated_sequence,