feat: add device

app.py

@@ -1,26 +1,18 @@
 import torch
-import argparse
 import selfies as sf
-from tqdm import tqdm
 from transformers import T5EncoderModel
-from transformers import set_seed
 from src.scripts.mytokenizers import Tokenizer
 from src.improved_diffusion import gaussian_diffusion as gd
-from src.improved_diffusion import dist_util, logger
 from src.improved_diffusion.respace import SpacedDiffusion
 from src.improved_diffusion.transformer_model import TransformerNetModel
-from src.improved_diffusion.script_util import (
-    model_and_diffusion_defaults,
-    add_dict_to_argparser,
-)
-from src.scripts.mydatasets import Lang2molDataset_submission
 import streamlit as st
 import os
 
 
 @st.cache_resource
-def get_encoder():
+def get_encoder(device):
     model = T5EncoderModel.from_pretrained("QizhiPei/biot5-base-text2mol")
+    model.to(device)
     model.eval()
     return model
 
@@ -31,7 +23,7 @@ def get_tokenizer():
 
 
 @st.cache_resource
-def get_model():
+def get_model(device):
     model = TransformerNetModel(
         in_channels=32,
         model_channels=128,
@@ -44,9 +36,10 @@ def get_model():
     model.load_state_dict(
         torch.load(
             os.path.join("checkpoints", "PLAIN_ema_0.9999_360000.pt"),
-            map_location=torch.device("cpu"),
+            map_location=torch.device(device),
         )
     )
+    model.to(device)
     model.eval()
     return model
 
@@ -65,9 +58,11 @@ def get_diffusion():
     )
 
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
 tokenizer = get_tokenizer()
-encoder = get_encoder()
-model = get_model()
+encoder = get_encoder(device)
+model = get_model(device)
 diffusion = get_diffusion()
 
 st.title("Lang2mol-Diff")
@@ -85,8 +80,8 @@ if button:
             return_attention_mask=True,
         )
         caption_state = encoder(
-            input_ids=output["input_ids"],
-            attention_mask=output["attention_mask"],
+            input_ids=output["input_ids"].to(device),
+            attention_mask=output["attention_mask"].to(device),
         ).last_hidden_state
         caption_mask = output["attention_mask"]
 
@@ -98,7 +93,7 @@ if button:
             model_kwargs={},
             top_p=1.0,
             progress=True,
-            caption=(caption_state, caption_mask),
+            caption=(caption_state.to(device), caption_mask.to(device)),
         )
         logits = model.get_logits(torch.tensor(outputs))
         cands = torch.topk(logits, k=1, dim=-1)