Update demo.py

demo.py

@@ -25,8 +25,6 @@ tqdm.__init__ = partialmethod(tqdm.__init__, leave=False)
 import streamlit as st
 from stmol import *
 
-def download_callback(newFileName):
-    st.success(f"{newFileName}: successfully downloaded ")
 def download(outputFile,newFileName,description):
     with open(outputFile, "rb") as file:
         btn = st.download_button(
@@ -40,7 +38,7 @@ def download(outputFile,newFileName,description):
 import pandas as pd
 def display(output,style,resn):
     # imformation
-    protein=Protein.from_PDB(output)
+    protein=Protein.from_PDB(output,device=device)
     st.subheader("Protein Information:")
     st.write(f"Device: {protein.device}")
     st.write(f"Protein Length: {len(protein)} residues")
@@ -49,7 +47,7 @@ def display(output,style,resn):
     # 显示 Protein 的序列
     st.subheader("Protein Sequence:")
     protein_sequence = protein.sequence(format="three-letter-list")
-    st.markdown(f"**My List:** {protein_sequence}")
+    st.markdown(f"**Protein Sequence:** {protein_sequence}")
     st.write(protein_sequence)
     # 显示 Protein 的结构
     with open(output, "r") as file:
@@ -90,8 +88,7 @@ from chroma.utility.chroma import letter_to_point_cloud, plane_split_protein
 
 register_key(api_key)
 
-# device = 'cuda:2' if torch.cuda.is_available() else 'cpu'
-device='cuda:0'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 with contextlib.redirect_stdout(None):
     chroma = Chroma(device=device)