Prevent Gradio startup crash via deferred DB loading

app.py

@@ -91,7 +91,7 @@ custom_css = """
         height:190px !important;
     }
     .gallery-height {
-        height: 380px !important;
+        height: 350px !important;
     }
     #custom_plot {
        height: 300px !important;
@@ -128,11 +128,25 @@ def search_structure_from_mass(structureDB,mass, ppm):
     structures = structureDB[(structureDB['MonoisotopicMass'] >= mmin) & (structureDB['MonoisotopicMass'] <= mmax)]
     return structures 
 
-dataset_repo = "Tingxie/CSU-MS2-DB"
-db_filename = "csu_ms2_db.db"
-token = os.getenv("HF_TOKEN")
-db_path = hf_hub_download(repo_id=dataset_repo, filename=db_filename, repo_type="dataset", token=token)
-conn = sqlite3.connect(db_path, check_same_thread=False)
+conn = None 
+
+def initialize_db():
+    global conn
+    if conn is None:
+        dataset_repo = "Tingxie/CSU-MS2-DB"
+        db_filename = "csu_ms2_db.db"
+        token = os.getenv("HF_TOKEN")
+        print("Starting large file download and DB connection...")
+        db_path = hf_hub_download(repo_id=dataset_repo, filename=db_filename, repo_type="dataset", token=token)
+        conn = sqlite3.connect(db_path, check_same_thread=False)
+        print("DB initialization complete.")
+    return conn
+    
+#dataset_repo = "Tingxie/CSU-MS2-DB"
+#db_filename = "csu_ms2_db.db"
+#token = os.getenv("HF_TOKEN")
+#db_path = hf_hub_download(repo_id=dataset_repo, filename=db_filename, repo_type="dataset", token=token)
+#conn = sqlite3.connect(db_path, check_same_thread=False)
 
 device='cpu'
 pretrain_model_path_low,pretrain_model_path_median,pretrain_model_path_high='model/low_energy/checkpoints/model.pth','model/median_energy/checkpoints/model.pth','model/high_energy/checkpoints/model.pth'
@@ -192,8 +206,8 @@ def MS2Embedding(spectra):
 def calculate_cosine_similarity(vector1, vector2):
     return cosine_similarity(vector1.reshape(1, -1), vector2.reshape(1, -1))[0][0]
 
-def retrieve_similarity_scores( table_name, target_mass,collision_energy, ms2_embedding_low, ms2_embedding_median, ms2_embedding_high):
-    cur = conn.cursor()
+def retrieve_similarity_scores( conn_obj, table_name, target_mass,collision_energy, ms2_embedding_low, ms2_embedding_median, ms2_embedding_high):
+    cur = conn_obj.cursor()
     if table_name == "CSU_MS2_DB":
         table_name = 'ConSSDB'
     if table_name == "BloodExp: blood exposome database":
@@ -312,14 +326,14 @@ def get_topK_result(library,ms_feature, smiles_feature, topK):
     return indices, scores, candidates
 
 
-def rank_lib(database_name,spectrum_path,instrument_type,adduct,parent_Mass,collision_energy):
+def rank_lib(conn_obj, database_name,spectrum_path,instrument_type,adduct,parent_Mass,collision_energy):
     ms2 = list(load_from_msp(spectrum_path.name))[0]
     ms2 = spectrum_processing(ms2)
     collision_energy=float(collision_energy)
     parent_Mass=float(parent_Mass)
     ms2_embedding_low,ms2_embedding_median,ms2_embedding_high = MS2Embedding(ms2)
     ms2_embedding_low,ms2_embedding_median,ms2_embedding_high = torch.tensor(ms2_embedding_low).float(),torch.tensor(ms2_embedding_median).float(),torch.tensor(ms2_embedding_high).float()
-    top_10_smiles = retrieve_similarity_scores(database_name,parent_Mass,collision_energy,ms2_embedding_low,ms2_embedding_median,ms2_embedding_high)
+    top_10_smiles = retrieve_similarity_scores(conn_obj, database_name,parent_Mass,collision_energy,ms2_embedding_low,ms2_embedding_median,ms2_embedding_high)
     smis = [x[0][0] for x in top_10_smiles]
     scores = [x[1] for x in top_10_smiles]
     images,image_descrips=[],[]
@@ -411,6 +425,7 @@ def rank_user_lib(candidate_file,spectrum_path,instrument_type,adduct,parent_Mas
 with gr.Blocks(theme=seafoam) as demo:
     gr.HTML(custom_css)
     gr.Markdown('<div style="font-size:50px; font-weight:bold;">🔍 CSU-MS2 web server </div>')
+    db_conn_state = gr.State(None)
     with gr.Row():
         with gr.Column():
             peak_data = gr.File(file_count="single", label="Upload MS/MS spectrum file in .msp format", elem_classes=".file-upload-height")
@@ -472,7 +487,8 @@ with gr.Blocks(theme=seafoam) as demo:
         with gr.Column():
             user_button = gr.Button("Cross-Modal Retrieval")
             user_output = gr.Gallery(height='auto',columns=4,elem_classes="gallery-height",label='Cross-modal retrieval results')
-    lib_button.click(rank_lib, inputs=[dataset,peak_data,instru,ionmode,par_ion_mass,collision_e], outputs=lib_output)
+    demo.load(fn=initialize_db, inputs=None, outputs=db_conn_state, queue=True, show_progress="full")
+    lib_button.click(rank_lib, inputs=[db_conn_state, dataset,peak_data,instru,ionmode,par_ion_mass,collision_e], outputs=lib_output)
     user_button.click(rank_user_lib, inputs=[user_dataset,peak_data,instru,ionmode,par_ion_mass,collision_e], outputs=user_output)
 demo.launch(share=True)