feat: use gpu space

app.py

@@ -6,6 +6,7 @@ from src.improved_diffusion import gaussian_diffusion as gd
 from src.improved_diffusion.respace import SpacedDiffusion
 from src.improved_diffusion.transformer_model import TransformerNetModel
 import streamlit as st
+import spaces
 import os
 
 
@@ -57,28 +58,18 @@ def get_diffusion():
         training_mode="e2e",
     )
 
-
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-
-tokenizer = get_tokenizer()
-encoder = get_encoder(device)
-model = get_model(device)
-diffusion = get_diffusion()
-
-st.title("Lang2mol-Diff")
-text_input = st.text_area("Enter molecule description")
-button = st.button("Submit")
-if button:
+@spaces.GPU
+def generate(text_input):
     with st.spinner("Please wait..."):
         output = tokenizer(
-            text_input,
-            max_length=256,
-            truncation=True,
-            padding="max_length",
-            add_special_tokens=True,
-            return_tensors="pt",
-            return_attention_mask=True,
-        )
+                text_input,
+                max_length=256,
+                truncation=True,
+                padding="max_length",
+                add_special_tokens=True,
+                return_tensors="pt",
+                return_attention_mask=True,
+            )
         caption_state = encoder(
             input_ids=output["input_ids"].to(device),
             attention_mask=output["attention_mask"].to(device),
@@ -103,5 +94,18 @@ if button:
         result = sf.decoder(
             outputs[0].replace("<pad>", "").replace("</s>", "").replace("\t", "")
         ).replace("\t", "")
+        return result
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+tokenizer = get_tokenizer()
+encoder = get_encoder(device)
+model = get_model(device)
+diffusion = get_diffusion()
 
-        st.write(result)
+st.title("Lang2mol-Diff")
+text_input = st.text_area("Enter molecule description")
+button = st.button("Submit")
+if button:
+    result = generate(text_input)
+    st.write(result)