Update app.py

app.py

@@ -503,8 +503,7 @@ def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, num_samples
     ]
 
 
-demo = gr.Blocks()
-with demo:
+with gr.Blocks(js=output.STARTUP_JS) as demo:
     gr.Markdown('# DiffLinker: Equivariant 3D-Conditional Diffusion Model for Molecular Linker Design')
     gr.Markdown(
         'Given a set of disconnected fragments in 3D, '
@@ -596,14 +595,13 @@ with demo:
         fn=generate,
         inputs=[input_fragments_file, input_protein_file, n_steps, n_atoms, n_samples, hidden],
         outputs=[visualization, output_files, samples, hidden],
-        _js=output.RETURN_SELECTION_JS,
+        js=output.RETURN_SELECTION_JS,
     )
     samples.select(
         fn=draw_sample,
         inputs=[samples, output_files, n_samples],
         outputs=[visualization],
     )
-    demo.load(_js=output.STARTUP_JS)
 
 demo.launch(server_name=args.ip)