Update app.py

app.py

@@ -39,15 +39,16 @@ def get_correct_qcs_image(srv, num_of_qubits, max_gates, g):
     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) 
     qc_list, _, svr_list = convert_tensors_to_srvs(out_tensor, pipeline.gate_pool)
 
-    fig, axs = plt.subplots(3, 2, figsize=(7,20), constrained_layout=True, dpi=120)
+    fig, axs = plt.subplots(3, 2, figsize=(7,10), constrained_layout=True, dpi=120)
     for qc,is_svr,ax in zip(qc_list, svr_list, axs.flatten()): 
         qc.draw("mpl", plot_barriers=False, ax=ax)
         ax.set_title(f"{'Correct' if is_svr==srv else 'NOT correct'}, is SRV = {is_svr}")
 
-    buf = io.BytesIO()
-    fig.savefig(buf)
-    buf.seek(0)
-    return Image.open(buf)
+    # buf = io.BytesIO()
+    # fig.savefig(buf)
+    # buf.seek(0)
+    # return Image.open(buf)
+    return fig
 
 #--------------------------------
 # run
@@ -65,5 +66,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'):    
-    image = get_correct_qcs_image(ast.literal_eval(srv), num_of_qubits, max_gates, g)
-    col2.image(image, use_column_width=True)
+    fig = get_correct_qcs_image(ast.literal_eval(srv), num_of_qubits, max_gates, g)
+    # col2.image(image, use_column_width=True)
+    col2.pyplot(fig)