Update app.py

app.py

@@ -26,14 +26,17 @@ weigths_file = download(url_weights, save_destination)
 #--------------------------------
 # setup
 
-try:
-    pipeline
-except:
+@st.cache_resource
+def load_pipeline():
     pipeline = DiffusionPipeline.from_config_file(save_destination, infer_torch_device())  
-    pipeline.scheduler.set_timesteps(20) 
+    pipeline.scheduler.set_timesteps(20)   
+    return pipeline
+
+pipeline = load_pipeline()
+
 
 is_gpu_busy = False
-def get_correct_qcs_image(srv, num_of_qubits, max_gates, g):
+def get_qcs(srv, num_of_qubits, max_gates, g):
     global is_gpu_busy
        
     out_tensor           = generate_srv_tensors(pipeline, f"Generate SRV: {srv}", samples=6, system_size=num_of_qubits, num_of_qubits=num_of_qubits, max_gates=max_gates, g=g) 
@@ -66,6 +69,6 @@ max_gates     = col1.select_slider('Max gates', options=[4,8,12,16,20,24,28], va
 g             = col1.slider('Guidance scale', min_value=0.0, max_value=15.0, value=7.5)
 
 if col1.button('Generate circuits'):    
-    fig = get_correct_qcs_image(ast.literal_eval(srv), num_of_qubits, max_gates, g)
+    fig = get_qcs(ast.literal_eval(srv), num_of_qubits, max_gates, g)
     # col2.image(image, use_column_width=True)
     col2.pyplot(fig)