Bug fixes

app/fn.py

@@ -1,3 +1,4 @@
+import urllib.parse
 from datetime import datetime
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
@@ -18,7 +19,7 @@ import panel as pn
 import seaborn as sns
 from markdown import markdown
 from rdkit import Chem, RDConfig
-from rdkit.Chem import Crippen, Descriptors, rdMolDescriptors, Lipinski, rdmolops
+from rdkit.Chem import Crippen, Descriptors, rdMolDescriptors, Lipinski, rdmolops, Draw, rdDepictor
 import requests
 
 from app import static
@@ -37,6 +38,16 @@ COL_ALIASES = {
     'name': 'Complex Name',
 }
 
+COL_DTYPE = {
+    'out_path': 'str',
+    'ligand_conf_path': 'str',
+    'ID1': 'str',
+    'ID2': 'str',
+    'X1': 'str',
+    'X1^': 'str',
+    'name': 'str',
+}
+
 
 def lipinski(mol):
     """
@@ -401,14 +412,46 @@ def read_molecule_file(in_file, allowed_extentions):
 #     html = static.COMPLEX_RENDERING_TEMPLATE.format(viewer_models=viewer_models)
 #     return static.IFRAME_TEMPLATE.format(html=html)
 
-def create_result_table_html(summary_df, result_info, opts=(), progress=gr.Progress(track_tqdm=True)):
+def prepare_df_for_table(result_df):
+    result_df.dropna(subset=['mol'], inplace=True)
+
+    rdDepictor.SetPreferCoordGen(True)
+    draw_opts = Draw.rdMolDraw2D.MolDrawOptions()
+    draw_opts.clearBackground = False
+    draw_opts.bondLineWidth = 0.5
+    draw_opts.explicitMethyl = True
+    draw_opts.singleColourWedgeBonds = True
+    draw_opts.addStereoAnnotation = False
+    draw_opts.useCDKAtomPalette()
+
+    def draw_mol(mol):
+        # Create a new drawer instance for each molecule (for efficiency)
+        drawer = Draw.MolDraw2DSVG(90, 56)
+        drawer.SetDrawOptions(draw_opts)
+
+        # Draw the molecule and return the SVG as a URI
+        drawer.DrawMolecule(mol)
+        drawer.FinishDrawing()
+        return urllib.parse.quote(drawer.GetDrawingText())
+
+    # Convert to URI-formatted inline SVG
+    result_df['Compound'] = result_df['mol'].apply(draw_mol)
+
+    return result_df
+
+
+def create_result_table_html(summary_df, result_info=None, opts=(), progress=gr.Progress(track_tqdm=True)):
     html_df = summary_df.copy().drop(columns=['mol'])
-    output_dir = Path(result_info['output_dir'])
-    job_type = result_info['type']
     html_df.rename(columns=COL_ALIASES, inplace=True)
-    html_df['Pose'] = html_df['Pose'].apply(lambda x: str(output_dir / job_type / x))
-    # drop remaining columns ending with '_path'
+    if result_info:
+        output_dir = Path(result_info['output_dir'])
+        job_type = result_info['type']
+        html_df['Pose'] = html_df['Pose'].apply(lambda x: str(output_dir / job_type / x))
     hidden_cols = [col for col in html_df.columns if col.endswith('_path')]
+    rightmost_cols = ['Complex Name', 'Fragment SMILES', 'Compound SMILES']
+    col_order = ([col for col in html_df.columns if col not in rightmost_cols] +
+                 [col for col in html_df.columns if col in rightmost_cols])
+    html_df = html_df[col_order]
     html_df.index.name = 'Index'
 
     # if 'Scaffold' in html_df.columns and 'Exclude Scaffold Graph' not in opts:
@@ -448,7 +491,7 @@ def create_result_table_html(summary_df, result_info, opts=(), progress=gr.Progr
 
     report_table = pn.widgets.Tabulator(
         html_df, formatters=formatters,
-        frozen_columns=['Pose', 'Compound ID', 'Compound'],
+        frozen_columns=['Index', 'Pose', 'Compound ID', 'Compound'],
         hidden_columns=hidden_cols,
         sizing_mode='stretch_both',
         disabled=True, selectable=False,
@@ -499,17 +542,24 @@ def create_result_table_html(summary_df, result_info, opts=(), progress=gr.Progr
     #     template.main.append(
     #         pn.Card(stats_pane, sizing_mode='stretch_width', title='Summary Statistics', margin=10)
     #     )
-    report = pn.Column(pn.Accordion(
-        (f'{job_type.title()} Results', report_table),
-        toggle=True, margin=5, active=[0]
-    ))
+    if result_info:
+        table_title = (f"{job_type.title()} Results "
+                       f"({'No Linkable Pairs' if job_type != 'linking' else 'Generated Molecules'})")
+        report = pn.Column(pn.Accordion(
+            (table_title, report_table),
+            toggle=True, margin=5, active=[0]
+        ))
+        aspect_ratio = '1.090 / 1'
+    else:
+        report = report_table
+        aspect_ratio = '1.618 / 1'
 
     with tempfile.TemporaryDirectory() as tmpdir:
         file = Path(tmpdir) / 'report.html'
         report.save(file)
         # iframe_html = static.IFRAME_LINK_TEMPLATE.format(src="gradio_api/file=" + str(file))
         html_str = file.read_text()  # .replace('\'', '\"')
-        iframe_html = static.IFRAME_TEMPLATE.format(srcdoc=html.escape(html_str), aspect_ratio='1.090 / 1')
+        iframe_html = static.IFRAME_TEMPLATE.format(srcdoc=html.escape(html_str), aspect_ratio=aspect_ratio)
         return iframe_html
 
 

app/main.py

@@ -1,3 +1,7 @@
+import asyncio
+import queue
+import threading
+
 import spaces
 
 import io
@@ -7,8 +11,7 @@ import uuid
 import zipfile
 from datetime import datetime
 from pathlib import Path
-from time import sleep, time
-import urllib.parse
+from time import time
 
 from dotenv import load_dotenv
 import torch
@@ -19,7 +22,6 @@ from gradio_rangeslider import RangeSlider
 from omegaconf import OmegaConf
 import pandas as pd
 from rdkit import Chem
-from rdkit.Chem import PandasTools, Draw, rdDepictor
 
 from inference import (read_fragment_library, process_fragment_library, extract_pockets,
                        dock_fragments, generate_linkers, select_fragment_pairs)
@@ -27,12 +29,35 @@ from app import static, fn, db
 
 
 load_dotenv()
-Path(tempfile.gettempdir(), 'gradio').mkdir(exist_ok=True)
-SERVER_DATA_DIR = os.getenv('DATA', 'results')
-gr.set_static_paths(paths=["data/", SERVER_DATA_DIR, "app/"])
+RESULTS_DIR = os.getenv('RESULTS', 'results')
+MAX_CONCURRENT_JOBS = 1
+
+task_queue = queue.Queue()
+semaphore = threading.Semaphore(MAX_CONCURRENT_JOBS)
+lock = threading.Lock()
 job_db = db.init_job_db()
+
+Path(tempfile.gettempdir(), 'gradio').mkdir(exist_ok=True)
+gr.set_static_paths(paths=["data/", RESULTS_DIR, "app/"])
 os.chmod('./fpocket', 0o755)
 
+def task_worker():
+    """Worker function to process tasks from the queue with concurrency limit."""
+    while True:
+        data = task_queue.get()  # Get the next task from the queue
+        with semaphore:  # Ensure only 'MAX_CONCURRENT_JOBS' tasks run at once
+            with lock:  # Ensure only one task is processed at a time (for shared state)
+                dock_link(*data)
+        task_queue.task_done()
+
+
+worker_threads = []
+for _ in range(MAX_CONCURRENT_JOBS):
+    worker_thread = threading.Thread(target=task_worker, daemon=True)
+    worker_threads.append(worker_thread)
+    worker_thread.start()
+
+
 FRAG_LIBS = {'': None} | {
     lib_path.stem.replace('_', ' '): str(lib_path) for lib_path in Path('data/fragment_libraries').glob('*')
 }
@@ -58,6 +83,7 @@ POCKET_EXTRACT_OPTS = {
     }
 }
 
+
 def gr_error_wrapper(func):
     def wrapper(*args, **kwargs):
         try:
@@ -67,69 +93,71 @@ def gr_error_wrapper(func):
 
     return wrapper
 
-def query_job_status(job_id):
-    gr.Info('Start querying the job database...')
+
+async def query_job_status(job_id):
     stop = False
+    interval = 3  # Check every 3 seconds for better responsiveness
     retry = 0
+
     while not stop:
-        try:
-            job = job_db.job_lookup(job_id)
-            if job:
-                if job['status'] == "RUNNING":
-                    yield {
-                        pred_lookup_status: f'''
-Your  job (ID: **{job['id']}**) started at **{job['start_time']}** and is **RUNNING...**
-
-It might take a few minutes up to a few hours depending on the input size and the queue status.
-You may keep the page open and wait for job completion, or close the page and revisit later to look up the job status 
-using the job id. You will also receive an email notification once the job is done.
-''',
-                        pred_lookup_btn: gr.Button(visible=False),
-                        pred_lookup_stop_btn: gr.Button(visible=True),
-                    }
-                if job['status'] == "COMPLETED":
-                    stop = True
-                    msg = f"Your GenFBDD job (ID: {job['id']}) has been **COMPLETED**"
-                    msg += f" at **{job['end_time']}**" if job.get('end_time') else ""
-                    msg += f" and the results will **EXPIRE** by **{job['expiry_time']}**." if job.get('expiry_time') else "."
-                    msg += f' Redirecting to the Results page...'
-
-                    gr.Info(msg)
-                    yield {
-                        pred_lookup_status: msg,
-                        tabs: gr.Tabs(selected='result'),
-                        result_state: job,
-                        pred_lookup_btn: gr.Button(visible=True),
-                        pred_lookup_stop_btn: gr.Button(visible=False),
-                    }
-                if job['status'] == "FAILED":
-                    stop = True
-                    msg = f'Your GenFBDD job (ID: {job_id}) has **FAILED**'
-                    msg += f" at {job['end_time']}" if job.get('end_time') else ''
-                    msg += f" due to error: {job['error']}." if job.get('expiry_time') else '.'
-                    gr.Info(msg)
-                    yield {
-                        pred_lookup_status: msg,
-                        pred_lookup_btn: gr.Button(visible=True),
-                        pred_lookup_stop_btn: gr.Button(visible=False),
-                    }
-            else:
-                stop = (retry > 2)
-                if not stop:
-                    msg = f'Job ID {job_id} not found. Retrying... ({retry})'
-                else:
-                    msg = f'Job ID {job_id} not found after {retry} retries. Please double-check the job ID.'
-                gr.Info(msg)
-                retry += 1
+        # Wait for a short interval before checking again
+        await asyncio.sleep(interval)  # Non-blocking sleep
+        job = job_db.job_lookup(job_id)
+
+        if job:  # If the job exists
+            if job['status'] == "RUNNING":  # If the job is still running
+                yield {
+                    pred_lookup_status: f'''
+        Your job (ID: **{job['id']}**) started at **{job['start_time']}** and is **RUNNING...**
+
+        It might take a few minutes to a few hours depending on the input size and the queue status.
+        You may keep the page open or close it and revisit later using the job ID.
+        You will receive an email notification once the job is done.
+        ''',
+                    pred_lookup_btn: gr.update(visible=False),
+                    pred_lookup_stop_btn: gr.update(visible=True),
+                }
+
+            elif job['status'] == "COMPLETED":  # If the job is complete
+                stop = True
+                msg = f"Your job (ID: {job['id']}) has been **COMPLETED**"
+                msg += f" at **{job['end_time']}**" if job.get('end_time') else ""
+                msg += f" and the results will **EXPIRE** by **{job['expiry_time']}**." if job.get('expiry_time') else "."
+                msg += " Redirecting to the results page..."
+
                 yield {
                     pred_lookup_status: msg,
-                    pred_lookup_btn: gr.Button(visible=stop),
-                    pred_lookup_stop_btn: gr.Button(visible= not stop),
+                    tabs: gr.Tabs(selected='result'),
+                    result_state: job,
+                    pred_lookup_btn: gr.update(visible=True),
+                    pred_lookup_stop_btn: gr.update(visible=False),
                 }
-            sleep(5)
 
-        except Exception as e:
-            raise gr.Error(f'Failed to retrieve job status due to error: {str(e)}')
+            elif job['status'] == "FAILED":  # If the job failed
+                stop = True
+                msg = f'Your job (ID: {job_id}) has **FAILED**'
+                msg += f" at {job['end_time']}" if job.get('end_time') else ''
+                msg += f" due to error: {job['error']}." if job.get('error') else '.'
+
+                yield {
+                    pred_lookup_status: msg,
+                    pred_lookup_btn: gr.update(visible=True),
+                    pred_lookup_stop_btn: gr.update(visible=False),
+                }
+
+        else:  # If the job is not found
+            stop = retry > 2  # Stop after 3 retries
+            if not stop:
+                msg = f'Job ID {job_id} not found. Retrying... ({retry})'
+            else:
+                msg = f'Job ID {job_id} not found after {retry} retries. Please double-check the job ID.'
+
+            retry += 1
+            yield {
+                pred_lookup_status: msg,
+                pred_lookup_btn: gr.update(visible=stop),
+                pred_lookup_stop_btn: gr.update(visible=not stop),
+            }
 
 
 def checkbox_group_selections_to_kwargs(selected_options, option_mapping):
@@ -140,10 +168,12 @@ def checkbox_group_selections_to_kwargs(selected_options, option_mapping):
     return kwargs
 
 
-def job_validate(
-        frag_file, frag_df, prot_file,
+def job_submit(
+        frag_df, frag_file, prot_file,
+        dock_n_steps, dock_n_poses, dock_confidence_threshold,
+        linker_frag_dist, linker_strategy, linker_n_mols, linker_size, linker_steps,
         pocket_name, pocket_method, pocket_fs,
-        email, run_info, session_info: gr.Request
+        email, session_info: gr.Request
 ):
     if len(frag_df) == 0 or not frag_file:
         raise gr.Error("Please provide a valid fragment library.")
@@ -168,14 +198,10 @@ def job_validate(
         except EmailNotValidError as e:
             raise gr.Error(f"Invalid email address: {str(e)}.")
 
-    if run_info:
-        raise gr.Error(f"You already have a running prediction job (ID: {run_info['id']}) under this session. "
-                       "Please wait for it to complete before submitting another job.")
     if check := job_db.check_user_running_job(email, session_info):
         raise gr.Error(check)
 
-    gr.Info('Finished processing inputs. Initiating the GenFBDD job... '
-            'You will be redirected to Job Status page.')
+    gr.Info('Finished processing inputs. Initiating the GenFBDD job... You will be redirected to Job Status page.')
     job_id = str(uuid.uuid4())
     job_info = {
         'id': job_id,
@@ -194,7 +220,17 @@ def job_validate(
     }
     job_db.insert(job_info)
 
-    return job_info
+    task_queue.put((
+        frag_df, prot_file,
+        dock_n_steps, dock_n_poses, dock_confidence_threshold,
+        linker_frag_dist, linker_strategy, linker_n_mols, linker_size, linker_steps,
+        job_info
+    ))
+
+    return {
+        pred_lookup_id: job_id,
+        tabs: gr.Tabs(selected='job'),
+    }
 
 @spaces.GPU(duration=600)
 def dock_link(
@@ -216,7 +252,7 @@ def dock_link(
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(f'Using device: {device}')
     date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    out_dir = Path(SERVER_DATA_DIR, f'{date_time}_{job_id}')
+    out_dir = Path(RESULTS_DIR, f'{date_time}_{job_id}')
     frag_lib['X2'] = prot
     frag_lib['ID2'] = str(Path(prot).stem)
 
@@ -280,7 +316,6 @@ def dock_link(
             'output_dir': str(out_dir),
             'type': job_type,
         }
-        # return {result_state: job_info | update_info, run_state: {}}
 
     except Exception as e:
         gr.Warning(f"Job failed due to error: {str(e)}")
@@ -289,7 +324,6 @@ def dock_link(
             'error': str(e),
             'output_dir': None
         }
-        # return {result_state: {}, run_state: {}}
 
     finally:
         job_db.job_update(
@@ -298,10 +332,6 @@ def dock_link(
         )
 
 
-def get_session_state(request: gr.Request):
-    return request
-
-
 THEME = gr.themes.Base(
     spacing_size="sm", text_size='md', font=gr.themes.GoogleFont("Roboto"),
     primary_hue='emerald', secondary_hue='emerald', neutral_hue='slate',
@@ -329,8 +359,6 @@ THEME = gr.themes.Base(
 )
 
 with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600, 48 * 3600)) as demo:
-    run_state = gr.State(value={})
-    session_state = gr.State(value={})
     with gr.Column(variant='panel'):
         with gr.Tabs() as tabs:
             with gr.Tab(label='Start', id='start'):
@@ -361,11 +389,11 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
                         )
                         frag_lib_orig_df = gr.State(value=pd.DataFrame(columns=['X1', 'ID1', 'mol']))
                         frag_lib_mod_df = gr.State(value=pd.DataFrame(columns=['X1', 'ID1', 'mol']))
-                        # TODO: Tabulator with gr.HTML() for fragment library preview
-                        frag_lib_view = gr.DataFrame(
-                            value=pd.DataFrame(columns=['X1', 'ID1']), elem_id='frag_lib_view',
-                            visible=True, interactive=False,
-                        )
+                        # frag_lib_view = gr.DataFrame(
+                        #     value=pd.DataFrame(columns=['X1', 'ID1']), elem_id='frag_lib_view',
+                        #     visible=True, interactive=False,
+                        # )
+                        frag_lib_view = gr.HTML(static.IFRAME_TEMPLATE.format(aspect_ratio='1.618 /1', srcdoc=''))
 
                         with gr.Group():
                             frag_lib_process_opts = gr.CheckboxGroup(
@@ -510,6 +538,27 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
                             value='Run GenFBDD', variant='primary', interactive=True,
                         )
 
+            with gr.Tab(label='Jobs', id='job'):
+                gr.Markdown('''            
+                To check the status of an in-progress or historical job using the job ID and retrieve the predictions 
+                if the job has completed. Note that predictions are only kept for 48 hours upon job completion.
+                
+                You will be redirected to `Results` for carrying out further analysis and 
+                generating the full report when the job is done. If the the query fails to respond, please wait for a 
+                few minutes and refresh the page to try again. 
+                ''')
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        loader_html = gr.HTML('<div class="loader first-frame"></div>', visible=False)
+                    with gr.Column(scale=4):
+                        pred_lookup_id = gr.Textbox(
+                            label='Input Your Job ID', placeholder='e.g., e9dfd149-3f5c-48a6-b797-c27d027611ac',
+                            info="Your job ID is a UUID4 string that you receive after submitting a job on the "
+                                 "page or in the email notification.")
+                        pred_lookup_btn = gr.Button(value='Query Status', variant='primary', visible=True)
+                        pred_lookup_stop_btn = gr.Button(value='Stop Tracking', variant='stop', visible=False)
+                        pred_lookup_status = gr.Markdown("**Job Status**", container=True)
+
             with gr.Tab(label='Results', id='result'):
                 # Results
                 result_state = gr.State(value={})
@@ -520,6 +569,7 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
                     with gr.Row():
                         scores = gr.CheckboxGroup(list(fn.SCORE_MAP.keys()), label='Compound Scores')
                         filters = gr.CheckboxGroup(list(fn.FILTER_MAP.keys()), label='Compound Filters')
+                        result_example = gr.Button('Example', elem_classes=['example'])
                     with gr.Row():
                         prop_clr_btn = gr.ClearButton(value='Clear Properties', interactive=False)
                         prop_calc_btn = gr.Button(value='Calculate Properties', interactive=False, variant='primary')
@@ -531,28 +581,6 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
                             result_file_btn = gr.Button(value='Create Result File', visible=False, variant='primary')
                             result_download_file = gr.File(label='Download Result File', visible=False)
 
-            with gr.Tab(label='Job Status', id='job'):
-                gr.Markdown('''            
-                To check the status of an in-progress or historical job using the job ID and retrieve the predictions 
-                if the job has completed. Note that predictions are only kept for 48 hours upon job completion.
-                
-                You will be redirected to `Results` for carrying out further analysis and 
-                generating the full report when the job is done. If the the query fails to respond, please wait for a 
-                few minutes and refresh the page to try again. 
-                ''')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        loader_html = gr.HTML('<div class="loader first-frame"></div>', visible=False)
-                    with gr.Column(scale=4):
-                        pred_lookup_id = gr.Textbox(
-                            label='Input Your Job ID', placeholder='e.g., e9dfd149-3f5c-48a6-b797-c27d027611ac',
-                            info="Your job ID is a UUID4 string that you receive after submitting a job on the "
-                                 "page or in the email notification.")
-                        pred_lookup_example = gr.Button('Example', elem_classes=['example'], scale=1)
-                        pred_lookup_btn = gr.Button(value='Query Status', variant='primary', visible=True)
-                        pred_lookup_stop_btn = gr.Button(value='Stop Tracking', variant='stop', visible=False)
-                        pred_lookup_status = gr.Markdown("**Job Status**", container=True)
-
     # Event handlers
     ## Start tab
     ### Fragment Library
@@ -573,7 +601,7 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
         inputs=[frag_lib_file],
         outputs=[frag_lib_orig_df],
     ).success(
-        fn=lambda df: [df, gr.DataFrame(df.drop(columns='mol'), visible=True)],
+        fn=lambda df: [df.copy(), fn.create_result_table_html(fn.prepare_df_for_table(df))],
         inputs=[frag_lib_orig_df],
         outputs=[frag_lib_mod_df, frag_lib_view],
     )
@@ -586,7 +614,7 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
             new_df:=process_fragment_library(
                 df, **checkbox_group_selections_to_kwargs(opts, FRAG_LIB_PROCESS_OPTS)
             ),
-            gr.DataFrame(new_df.drop(columns='mol'), visible=True)
+            fn.create_result_table_html(fn.prepare_df_for_table(new_df))
         ],
         inputs=[frag_lib_orig_df, frag_lib_process_opts],
         outputs=[frag_lib_mod_df, frag_lib_view],
@@ -705,26 +733,16 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
         inputs=[selected_ligand, selected_pocket],
         outputs=[selected_ligand, selected_pocket],
     ).success(
-        fn=job_validate,
-        inputs=[
-            frag_lib_file, frag_lib_mod_df, input_prot_file,
-            selected_pocket, pocket_extract_dropdown, pocket_files,
-            email_input, run_state
-        ],
-        outputs=[run_state],
-    )
-
-    job_valid.success(
-        fn=dock_link,
+        fn=job_submit,
         inputs=[
-            frag_lib_mod_df, input_prot_file,
+            frag_lib_mod_df, frag_lib_file, input_prot_file,
             dock_steps, dock_n_poses, dock_confidence_cutoff,
             link_frag_dist_range, link_frag_pose_strategy, link_n_mols,
             link_linker_size, link_steps,
-            run_state
+            selected_pocket, pocket_extract_dropdown, pocket_files,
+            email_input,
         ],
-        # outputs=[result_state, run_state],
-        concurrency_limit=1, concurrency_id="gpu_queue"
+        outputs=[pred_lookup_id, tabs],
     )
 
     start_reset_components=[
@@ -750,57 +768,49 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
     )
 
     ### Job Status
-    pred_lookup_btn.click(
-        fn=lambda: '<div class="loader"></ div>',
-        outputs=loader_html,
-    )
     user_job_lookup = pred_lookup_btn.click(
+        fn=lambda: [
+            gr.update(value="Start querying the job database..."),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        ],
+        outputs=[pred_lookup_status, pred_lookup_btn, pred_lookup_stop_btn],
+    ).success(
         fn=query_job_status,
         inputs=[pred_lookup_id],
         outputs=[pred_lookup_status, tabs, result_state, pred_lookup_btn, pred_lookup_stop_btn],
         show_progress='minimal',
-    ).success(
-        lambda: '<div class="loader first-frame"></ div>',
-        outputs=loader_html,
     )
 
-    job_valid_success = job_valid.success(
-        fn=lambda job: [job['id'], gr.Tabs(selected='job')],
-        inputs=[run_state],
-        outputs=[pred_lookup_id, tabs],
+    auto_job_lookup = job_valid.success(
+        fn=lambda: [
+            gr.update(value="Start querying the job database..."),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        ],
+        outputs=[pred_lookup_status, pred_lookup_btn, pred_lookup_stop_btn],
     ).success(
-        lambda: '<div class="loader"></ div>',
-        outputs=loader_html,
-    )
-
-    auto_job_lookup = job_valid_success.success(
         fn=query_job_status,
         inputs=pred_lookup_id,
         outputs=[pred_lookup_status, tabs, result_state, pred_lookup_btn, pred_lookup_stop_btn],
         show_progress='minimal',
         cancels=[user_job_lookup],
-    ).success(
-        lambda: '<div class="loader first-frame"></ div>',
-        outputs=loader_html,
     )
 
     pred_lookup_stop_btn.click(
         fn=lambda: [gr.Button(visible=True), gr.Button(visible=False)],
         outputs=[pred_lookup_btn, pred_lookup_stop_btn],
         cancels=[user_job_lookup, auto_job_lookup],
-    ).success(
-        lambda: '<div class="loader first-frame"></ div>',
-        outputs=loader_html,
     )
 
-    pred_lookup_example.click(
+    result_example.click(
         fn=lambda: '80cf2658-7a1c-48d6-8372-61b978177fe6',
         outputs=[pred_lookup_id],
         show_progress='hidden'
     ).success(
         fn=query_job_status,
         inputs=pred_lookup_id,
-        outputs=[pred_lookup_status, tabs, result_state],
+        outputs=[pred_lookup_status, tabs, result_state, pred_lookup_btn, pred_lookup_stop_btn],
         show_progress='minimal',
         cancels=[user_job_lookup, auto_job_lookup],
     )
@@ -811,29 +821,16 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
         result_type = result_info['type']
         protein_structure_file = Path(result_info['protein_structure_file'])
         if result_type == 'docking':
-            result_df = pd.read_csv(result_dir / 'docking_summary.csv')
+            result_df = pd.read_csv(result_dir / 'docking_summary.csv', dtype=fn.COL_DTYPE)
             result_df['mol'] = result_df['X1'].apply(Chem.MolFromSmiles)
         elif result_type == 'linking':
-            result_df = pd.read_csv(result_dir / 'linking_summary.csv')
-            result_df = result_df[~result_df['X1^'].str.contains('.', regex=False)]
+            result_df = pd.read_csv(result_dir / 'linking_summary.csv', dtype=fn.COL_DTYPE)
             result_df['mol'] = result_df['X1^'].apply(Chem.MolFromSmiles)
-            result_df.dropna(subset=['mol'], inplace=True)
+            result_df = result_df[~result_df['X1^'].str.contains('.', regex=False)]
         else:
             raise gr.Error('Invalid result type')
 
-        rdDepictor.SetPreferCoordGen(True)
-        draw_opts = Draw.rdMolDraw2D.MolDrawOptions()
-        draw_opts.clearBackground = False
-        draw_opts.bondLineWidth = 0.5
-        draw_opts.explicitMethyl = True
-        draw_opts.singleColourWedgeBonds = True
-        draw_opts.addStereoAnnotation = False
-        draw_opts.useCDKAtomPalette()
-        PandasTools.drawOptions = draw_opts
-        PandasTools.molSize = (90, 56)
-        PandasTools.molRepresentation = 'svg'
-        # Convert to URI-formatted inline SVG
-        result_df['Compound'] = result_df['mol'].apply(PandasTools.PrintAsImageString).apply(urllib.parse.quote)
+        result_df = fn.prepare_df_for_table(result_df)
 
         return {
             result_table_orig_df: result_df,
@@ -921,13 +918,14 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
         outputs=[result_download_file],
     )
 
-    demo.load(fn=get_session_state, inputs=None, outputs=session_state, js=static.SETUP_JS)
+    demo.load(None, None, None, js=static.SETUP_JS)
 
+demo.share_token = 'genfbdd'
+demo.queue(default_concurrency_limit=None)
 demo.launch(
     server_name='0.0.0.0',
     max_file_size="5mb",
     ssr_mode=False,
-    share=True,
-    share_server_address="ciddr-lab.ac.cn:7000",
-    share_server_protocol="https"
+    show_api=False,
+    enable_monitoring=True,
 )

app/static.py

@@ -517,26 +517,41 @@ CREATE_OUTPUT_MOL_VIEW = """
             viewer.zoomTo();
             viewer.render();
 
-            const container = document.createElement("div");
-            container.classList.add("gr-btn-grp");
+            let container = element.parentElement.querySelector(".gr-btn-grp");
+            if (!container) {
+                container = document.createElement("div");
+                container.classList.add("gr-btn-grp");
+                element.parentElement.appendChild(container);
+            }
 
             // Molecule Button
-            const toggleMoleculeButton = document.createElement("button");
-            toggleMoleculeButton.textContent = "Hide Generated Molecule";
-            let moleculeIsHidden = false;
-            toggleMoleculeButton.onclick = function() {
-                if (viewer.models.length === 2) {
-                    moleculeIsHidden = !moleculeIsHidden;
-                    viewer.addStyle({model: 1}, {stick: {hidden: moleculeIsHidden} });
-                    toggleMoleculeButton.textContent = moleculeIsHidden ? "Show Generated Molecule" : "Hide Generated Molecule";
-                    viewer.render();
-                }
-            };
-            container.appendChild(toggleMoleculeButton);
+            let toggleMoleculeButton = container.querySelector("#toggleMoleculeButton");
+            if (!toggleMoleculeButton) {
+                toggleMoleculeButton = document.createElement("button");
+                toggleMoleculeButton.id = "toggleMoleculeButton"; // Add an ID
+                toggleMoleculeButton.textContent = "Hide Generated Molecule";
+                let moleculeIsHidden = false;
+                toggleMoleculeButton.onclick = function() {
+                    if (viewer.models.length === 2) {
+                        moleculeIsHidden = !moleculeIsHidden;
+                        viewer.addStyle({model: 1}, {stick: {hidden: moleculeIsHidden} });
+                        toggleMoleculeButton.textContent = moleculeIsHidden ? "Show Generated Molecule" : "Hide Generated Molecule";
+                        viewer.render();
+                    }
+                };
+                container.appendChild(toggleMoleculeButton);
+            }
 
             // Ligand Button
-            if (viewer.getAtomsFromSel({model: 0, hetflag : true}).length > 0) {
-                const toggleLigandButton = document.createElement("button");
+            const ligandButtonId = "toggleLigandButton";
+            let toggleLigandButton = container.querySelector("#" + ligandButtonId);
+
+            // Check for ligands and existing button
+            const hasLigands = viewer.getAtomsFromSel({model: 0, hetflag : true}).length > 0;
+            if (hasLigands && !toggleLigandButton) {
+                // Create button if ligands exist and it doesn't exist yet
+                toggleLigandButton = document.createElement("button");
+                toggleLigandButton.id = ligandButtonId;
                 toggleLigandButton.textContent = "Hide Co-Crystallized Ligand";
                 let ligandIsHidden = false;
                 toggleLigandButton.onclick = function() {
@@ -546,20 +561,27 @@ CREATE_OUTPUT_MOL_VIEW = """
                     viewer.render();
                 };
                 container.appendChild(toggleLigandButton);
+            } else if (!hasLigands && toggleLigandButton) {
+                // Remove button if no ligands exist and it already exists
+                container.removeChild(toggleLigandButton);
             }
 
             // Protein Button
-            const toggleProteinButton = document.createElement("button");
-            toggleProteinButton.textContent = "Hide Target Protein";
-            let proteinIsHidden = false;
-            toggleProteinButton.onclick = function() {
-                proteinIsHidden = !proteinIsHidden;
-                viewer.addStyle({model: 0, hetflag: false}, {cartoon: {hidden: proteinIsHidden} });
-                toggleProteinButton.textContent = proteinIsHidden ? "Show Target Protein" : "Hide Target Protein";
-                viewer.render();
-            };
-            container.appendChild(toggleProteinButton);
-            element.parentElement.appendChild(container);
+            let toggleProteinButton = container.querySelector("#toggleProteinButton");
+            if (!toggleProteinButton) {
+                toggleProteinButton = document.createElement("button");
+                toggleProteinButton.id = "toggleProteinButton";
+                toggleProteinButton.textContent = "Hide Target Protein";
+                let proteinIsHidden = false;
+                toggleProteinButton.onclick = function() {
+                    proteinIsHidden = !proteinIsHidden;
+                    viewer.addStyle({model: 0, hetflag: false}, {cartoon: {hidden: proteinIsHidden} });
+                    toggleProteinButton.textContent = proteinIsHidden ? "Show Target Protein" : "Hide Target Protein";
+                    viewer.render();
+                };
+                container.appendChild(toggleProteinButton);
+            }
+
         }).fail(function(error) {
             console.error("Error loading molecule:", error);
         });

app/static/medium-zoom.min.js

@@ -1,2 +0,0 @@
-/*! medium-zoom 1.1.0 | MIT License | https://github.com/francoischalifour/medium-zoom */
-!function(e,t){"object"==typeof exports&&"undefined"!=typeof module?module.exports=t():"function"==typeof define&&define.amd?define(t):(e=e||self).mediumZoom=t()}(this,(function(){"use strict";var e=Object.assign||function(e){for(var t=1;t<arguments.length;t++){var o=arguments[t];for(var n in o)Object.prototype.hasOwnProperty.call(o,n)&&(e[n]=o[n])}return e},t=function(e){return"IMG"===e.tagName},o=function(e){return e&&1===e.nodeType},n=function(e){return".svg"===(e.currentSrc||e.src).substr(-4).toLowerCase()},i=function(e){try{return Array.isArray(e)?e.filter(t):function(e){return NodeList.prototype.isPrototypeOf(e)}(e)?[].slice.call(e).filter(t):o(e)?[e].filter(t):"string"==typeof e?[].slice.call(document.querySelectorAll(e)).filter(t):[]}catch(e){throw new TypeError("The provided selector is invalid.\nExpects a CSS selector, a Node element, a NodeList or an array.\nSee: https://github.com/francoischalifour/medium-zoom")}},r=function(e){var t=document.createElement("div");return t.classList.add("medium-zoom-overlay"),t.style.background=e,t},d=function(e){var t=e.getBoundingClientRect(),o=t.top,n=t.left,i=t.width,r=t.height,d=e.cloneNode(),a=window.pageYOffset||document.documentElement.scrollTop||document.body.scrollTop||0,m=window.pageXOffset||document.documentElement.scrollLeft||document.body.scrollLeft||0;return d.removeAttribute("id"),d.style.position="absolute",d.style.top=o+a+"px",d.style.left=n+m+"px",d.style.width=i+"px",d.style.height=r+"px",d.style.transform="",d},a=function(t,o){var n=e({bubbles:!1,cancelable:!1,detail:void 0},o);if("function"==typeof window.CustomEvent)return new CustomEvent(t,n);var i=document.createEvent("CustomEvent");return i.initCustomEvent(t,n.bubbles,n.cancelable,n.detail),i};return function(e,t){void 0===t&&(t={});var o=t.insertAt;if(e&&"undefined"!=typeof document){var n=document.head||document.getElementsByTagName("head")[0],i=document.createElement("style");i.type="text/css","top"===o&&n.firstChild?n.insertBefore(i,n.firstChild):n.appendChild(i),i.styleSheet?i.styleSheet.cssText=e:i.appendChild(document.createTextNode(e))}}(".medium-zoom-overlay{position:fixed;top:0;right:0;bottom:0;left:0;opacity:0;transition:opacity .3s;will-change:opacity}.medium-zoom--opened .medium-zoom-overlay{cursor:pointer;cursor:zoom-out;opacity:1}.medium-zoom-image{cursor:pointer;cursor:zoom-in;transition:transform .3s cubic-bezier(.2,0,.2,1)!important}.medium-zoom-image--hidden{visibility:hidden}.medium-zoom-image--opened{position:relative;cursor:pointer;cursor:zoom-out;will-change:transform}"),function t(m){var l=arguments.length>1&&void 0!==arguments[1]?arguments[1]:{},c=window.Promise||function(e){function t(){}e(t,t)},u=function(e){var t=e.target;t!==N?-1!==x.indexOf(t)&&w({target:t}):E()},s=function(){if(!A&&k.original){var e=window.pageYOffset||document.documentElement.scrollTop||document.body.scrollTop||0;Math.abs(S-e)>T.scrollOffset&&setTimeout(E,150)}},f=function(e){var t=e.key||e.keyCode;"Escape"!==t&&"Esc"!==t&&27!==t||E()},p=function(){var t=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{},n=t;if(t.background&&(N.style.background=t.background),t.container&&t.container instanceof Object&&(n.container=e({},T.container,t.container)),t.template){var i=o(t.template)?t.template:document.querySelector(t.template);n.template=i}return T=e({},T,n),x.forEach((function(e){e.dispatchEvent(a("medium-zoom:update",{detail:{zoom:j}}))})),j},g=function(){var o=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{};return t(e({},T,o))},v=function(){for(var e=arguments.length,t=Array(e),o=0;o<e;o++)t[o]=arguments[o];var n=t.reduce((function(e,t){return[].concat(e,i(t))}),[]);return n.filter((function(e){return-1===x.indexOf(e)})).forEach((function(e){x.push(e),e.classList.add("medium-zoom-image")})),O.forEach((function(e){var t=e.type,o=e.listener,i=e.options;n.forEach((function(e){e.addEventListener(t,o,i)}))})),j},h=function(){for(var e=arguments.length,t=Array(e),o=0;o<e;o++)t[o]=arguments[o];k.zoomed&&E();var n=t.length>0?t.reduce((function(e,t){return[].concat(e,i(t))}),[]):x;return n.forEach((function(e){e.classList.remove("medium-zoom-image"),e.dispatchEvent(a("medium-zoom:detach",{detail:{zoom:j}}))})),x=x.filter((function(e){return-1===n.indexOf(e)})),j},z=function(e,t){var o=arguments.length>2&&void 0!==arguments[2]?arguments[2]:{};return x.forEach((function(n){n.addEventListener("medium-zoom:"+e,t,o)})),O.push({type:"medium-zoom:"+e,listener:t,options:o}),j},y=function(e,t){var o=arguments.length>2&&void 0!==arguments[2]?arguments[2]:{};return x.forEach((function(n){n.removeEventListener("medium-zoom:"+e,t,o)})),O=O.filter((function(o){return!(o.type==="medium-zoom:"+e&&o.listener.toString()===t.toString())})),j},b=function(){var t=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{},i=t.target,r=function(){var t={width:document.documentElement.clientWidth,height:document.documentElement.clientHeight,left:0,top:0,right:0,bottom:0},i=void 0,r=void 0;if(T.container)if(T.container instanceof Object)i=(t=e({},t,T.container)).width-t.left-t.right-2*T.margin,r=t.height-t.top-t.bottom-2*T.margin;else{var d=(o(T.container)?T.container:document.querySelector(T.container)).getBoundingClientRect(),a=d.width,m=d.height,l=d.left,c=d.top;t=e({},t,{width:a,height:m,left:l,top:c})}i=i||t.width-2*T.margin,r=r||t.height-2*T.margin;var u=k.zoomedHd||k.original,s=n(u)?i:u.naturalWidth||i,f=n(u)?r:u.naturalHeight||r,p=u.getBoundingClientRect(),g=p.top,v=p.left,h=p.width,z=p.height,y=Math.min(Math.max(h,s),i)/h,b=Math.min(Math.max(z,f),r)/z,E=Math.min(y,b),w="scale("+E+") translate3d("+((i-h)/2-v+T.margin+t.left)/E+"px, "+((r-z)/2-g+T.margin+t.top)/E+"px, 0)";k.zoomed.style.transform=w,k.zoomedHd&&(k.zoomedHd.style.transform=w)};return new c((function(e){if(i&&-1===x.indexOf(i))e(j);else{if(k.zoomed)e(j);else{if(i)k.original=i;else{if(!(x.length>0))return void e(j);var t=x;k.original=t[0]}if(k.original.dispatchEvent(a("medium-zoom:open",{detail:{zoom:j}})),S=window.pageYOffset||document.documentElement.scrollTop||document.body.scrollTop||0,A=!0,k.zoomed=d(k.original),document.body.appendChild(N),T.template){var n=o(T.template)?T.template:document.querySelector(T.template);k.template=document.createElement("div"),k.template.appendChild(n.content.cloneNode(!0)),document.body.appendChild(k.template)}if(k.original.parentElement&&"PICTURE"===k.original.parentElement.tagName&&k.original.currentSrc&&(k.zoomed.src=k.original.currentSrc),document.body.appendChild(k.zoomed),window.requestAnimationFrame((function(){document.body.classList.add("medium-zoom--opened")})),k.original.classList.add("medium-zoom-image--hidden"),k.zoomed.classList.add("medium-zoom-image--opened"),k.zoomed.addEventListener("click",E),k.zoomed.addEventListener("transitionend",(function t(){A=!1,k.zoomed.removeEventListener("transitionend",t),k.original.dispatchEvent(a("medium-zoom:opened",{detail:{zoom:j}})),e(j)})),k.original.getAttribute("data-zoom-src")){k.zoomedHd=k.zoomed.cloneNode(),k.zoomedHd.removeAttribute("srcset"),k.zoomedHd.removeAttribute("sizes"),k.zoomedHd.removeAttribute("loading"),k.zoomedHd.src=k.zoomed.getAttribute("data-zoom-src"),k.zoomedHd.onerror=function(){clearInterval(m),console.warn("Unable to reach the zoom image target "+k.zoomedHd.src),k.zoomedHd=null,r()};var m=setInterval((function(){k.zoomedHd.complete&&(clearInterval(m),k.zoomedHd.classList.add("medium-zoom-image--opened"),k.zoomedHd.addEventListener("click",E),document.body.appendChild(k.zoomedHd),r())}),10)}else if(k.original.hasAttribute("srcset")){k.zoomedHd=k.zoomed.cloneNode(),k.zoomedHd.removeAttribute("sizes"),k.zoomedHd.removeAttribute("loading");var l=k.zoomedHd.addEventListener("load",(function(){k.zoomedHd.removeEventListener("load",l),k.zoomedHd.classList.add("medium-zoom-image--opened"),k.zoomedHd.addEventListener("click",E),document.body.appendChild(k.zoomedHd),r()}))}else r()}}}))},E=function(){return new c((function(e){if(!A&&k.original){A=!0,document.body.classList.remove("medium-zoom--opened"),k.zoomed.style.transform="",k.zoomedHd&&(k.zoomedHd.style.transform=""),k.template&&(k.template.style.transition="opacity 150ms",k.template.style.opacity=0),k.original.dispatchEvent(a("medium-zoom:close",{detail:{zoom:j}})),k.zoomed.addEventListener("transitionend",(function t(){k.original.classList.remove("medium-zoom-image--hidden"),document.body.removeChild(k.zoomed),k.zoomedHd&&document.body.removeChild(k.zoomedHd),document.body.removeChild(N),k.zoomed.classList.remove("medium-zoom-image--opened"),k.template&&document.body.removeChild(k.template),A=!1,k.zoomed.removeEventListener("transitionend",t),k.original.dispatchEvent(a("medium-zoom:closed",{detail:{zoom:j}})),k.original=null,k.zoomed=null,k.zoomedHd=null,k.template=null,e(j)}))}else e(j)}))},w=function(){var e=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{},t=e.target;return k.original?E():b({target:t})},L=function(){return T},H=function(){return x},C=function(){return k.original},x=[],O=[],A=!1,S=0,T=l,k={original:null,zoomed:null,zoomedHd:null,template:null};"[object Object]"===Object.prototype.toString.call(m)?T=m:(m||"string"==typeof m)&&v(m),T=e({margin:0,background:"#fff",scrollOffset:40,container:null,template:null},T);var N=r(T.background);document.addEventListener("click",u),document.addEventListener("keyup",f),document.addEventListener("scroll",s),window.addEventListener("resize",E);var j={open:b,close:E,toggle:w,update:p,clone:g,attach:v,detach:h,on:z,off:y,getOptions:L,getImages:H,getZoomedImage:C};return j}}));

app/static/zooming.min.js

@@ -1 +0,0 @@
-!function(t,e){"object"==typeof exports&&"undefined"!=typeof module?module.exports=e():"function"==typeof define&&define.amd?define(e):(t=t||self).Zooming=e()}(this,function(){"use strict";var t="auto",e="zoom-in",i="zoom-out",n="grab",s="move";function o(t,e,i){var n={passive:!1};!(arguments.length>3&&void 0!==arguments[3])||arguments[3]?t.addEventListener(e,i,n):t.removeEventListener(e,i,n)}function r(t,e){if(t){var i=new Image;i.onload=function(){e&&e(i)},i.src=t}}function a(t){return t.dataset.original?t.dataset.original:"A"===t.parentNode.tagName?t.parentNode.getAttribute("href"):null}function l(t,e,i){!function(t){var e=h.transitionProp,i=h.transformProp;if(t.transition){var n=t.transition;delete t.transition,t[e]=n}if(t.transform){var s=t.transform;delete t.transform,t[i]=s}}(e);var n=t.style,s={};for(var o in e)i&&(s[o]=n[o]||""),n[o]=e[o];return s}var h={transitionProp:"transition",transEndEvent:"transitionend",transformProp:"transform",transformCssProp:"transform"},c=h.transformCssProp,u=h.transEndEvent;var d=function(){},f={enableGrab:!0,preloadImage:!1,closeOnWindowResize:!0,transitionDuration:.4,transitionTimingFunction:"cubic-bezier(0.4, 0, 0, 1)",bgColor:"rgb(255, 255, 255)",bgOpacity:1,scaleBase:1,scaleExtra:.5,scrollThreshold:40,zIndex:998,customSize:null,onOpen:d,onClose:d,onGrab:d,onMove:d,onRelease:d,onBeforeOpen:d,onBeforeClose:d,onBeforeGrab:d,onBeforeRelease:d,onImageLoading:d,onImageLoaded:d},p={init:function(t){var e,i;e=this,i=t,Object.getOwnPropertyNames(Object.getPrototypeOf(e)).forEach(function(t){e[t]=e[t].bind(i)})},click:function(t){if(t.preventDefault(),m(t))return window.open(this.target.srcOriginal||t.currentTarget.src,"_blank");this.shown?this.released?this.close():this.release():this.open(t.currentTarget)},scroll:function(){var t=document.documentElement||document.body.parentNode||document.body,e=window.pageXOffset||t.scrollLeft,i=window.pageYOffset||t.scrollTop;null===this.lastScrollPosition&&(this.lastScrollPosition={x:e,y:i});var n=this.lastScrollPosition.x-e,s=this.lastScrollPosition.y-i,o=this.options.scrollThreshold;(Math.abs(s)>=o||Math.abs(n)>=o)&&(this.lastScrollPosition=null,this.close())},keydown:function(t){(function(t){return"Escape"===(t.key||t.code)||27===t.keyCode})(t)&&(this.released?this.close():this.release(this.close))},mousedown:function(t){if(y(t)&&!m(t)){t.preventDefault();var e=t.clientX,i=t.clientY;this.pressTimer=setTimeout(function(){this.grab(e,i)}.bind(this),200)}},mousemove:function(t){this.released||this.move(t.clientX,t.clientY)},mouseup:function(t){y(t)&&!m(t)&&(clearTimeout(this.pressTimer),this.released?this.close():this.release())},touchstart:function(t){t.preventDefault();var e=t.touches[0],i=e.clientX,n=e.clientY;this.pressTimer=setTimeout(function(){this.grab(i,n)}.bind(this),200)},touchmove:function(t){if(!this.released){var e=t.touches[0],i=e.clientX,n=e.clientY;this.move(i,n)}},touchend:function(t){(function(t){t.targetTouches.length})(t)||(clearTimeout(this.pressTimer),this.released?this.close():this.release())},clickOverlay:function(){this.close()},resizeWindow:function(){this.close()}};function y(t){return 0===t.button}function m(t){return t.metaKey||t.ctrlKey}var g={init:function(t){this.el=document.createElement("div"),this.instance=t,this.parent=document.body,l(this.el,{position:"fixed",top:0,left:0,right:0,bottom:0,opacity:0}),this.updateStyle(t.options),o(this.el,"click",t.handler.clickOverlay.bind(t))},updateStyle:function(t){l(this.el,{zIndex:t.zIndex,backgroundColor:t.bgColor,transition:"opacity\n        "+t.transitionDuration+"s\n        "+t.transitionTimingFunction})},insert:function(){this.parent.appendChild(this.el)},remove:function(){this.parent.removeChild(this.el)},fadeIn:function(){this.el.offsetWidth,this.el.style.opacity=this.instance.options.bgOpacity},fadeOut:function(){this.el.style.opacity=0}},v="function"==typeof Symbol&&"symbol"==typeof Symbol.iterator?function(t){return typeof t}:function(t){return t&&"function"==typeof Symbol&&t.constructor===Symbol&&t!==Symbol.prototype?"symbol":typeof t},b=function(t,e){if(!(t instanceof e))throw new TypeError("Cannot call a class as a function")},w=function(){function t(t,e){for(var i=0;i<e.length;i++){var n=e[i];n.enumerable=n.enumerable||!1,n.configurable=!0,"value"in n&&(n.writable=!0),Object.defineProperty(t,n.key,n)}}return function(e,i,n){return i&&t(e.prototype,i),n&&t(e,n),e}}(),x=Object.assign||function(t){for(var e=1;e<arguments.length;e++){var i=arguments[e];for(var n in i)Object.prototype.hasOwnProperty.call(i,n)&&(t[n]=i[n])}return t},O={init:function(t,e){this.el=t,this.instance=e,this.srcThumbnail=this.el.getAttribute("src"),this.srcset=this.el.getAttribute("srcset"),this.srcOriginal=a(this.el),this.rect=this.el.getBoundingClientRect(),this.translate=null,this.scale=null,this.styleOpen=null,this.styleClose=null},zoomIn:function(){var t=this.instance.options,e=t.zIndex,s=t.enableGrab,o=t.transitionDuration,r=t.transitionTimingFunction;this.translate=this.calculateTranslate(),this.scale=this.calculateScale(),this.styleOpen={position:"relative",zIndex:e+1,cursor:s?n:i,transition:c+"\n        "+o+"s\n        "+r,transform:"translate3d("+this.translate.x+"px, "+this.translate.y+"px, 0px)\n        scale("+this.scale.x+","+this.scale.y+")",height:this.rect.height+"px",width:this.rect.width+"px"},this.el.offsetWidth,this.styleClose=l(this.el,this.styleOpen,!0)},zoomOut:function(){this.el.offsetWidth,l(this.el,{transform:"none"})},grab:function(t,e,i){var n=k(),o=n.x-t,r=n.y-e;l(this.el,{cursor:s,transform:"translate3d(\n        "+(this.translate.x+o)+"px, "+(this.translate.y+r)+"px, 0px)\n        scale("+(this.scale.x+i)+","+(this.scale.y+i)+")"})},move:function(t,e,i){var n=k(),s=n.x-t,o=n.y-e;l(this.el,{transition:c,transform:"translate3d(\n        "+(this.translate.x+s)+"px, "+(this.translate.y+o)+"px, 0px)\n        scale("+(this.scale.x+i)+","+(this.scale.y+i)+")"})},restoreCloseStyle:function(){l(this.el,this.styleClose)},restoreOpenStyle:function(){l(this.el,this.styleOpen)},upgradeSource:function(){if(this.srcOriginal){var t=this.el.parentNode;this.srcset&&this.el.removeAttribute("srcset");var e=this.el.cloneNode(!1);e.setAttribute("src",this.srcOriginal),e.style.position="fixed",e.style.visibility="hidden",t.appendChild(e),setTimeout(function(){this.el.setAttribute("src",this.srcOriginal),t.removeChild(e)}.bind(this),50)}},downgradeSource:function(){this.srcOriginal&&(this.srcset&&this.el.setAttribute("srcset",this.srcset),this.el.setAttribute("src",this.srcThumbnail))},calculateTranslate:function(){var t=k(),e=this.rect.left+this.rect.width/2,i=this.rect.top+this.rect.height/2;return{x:t.x-e,y:t.y-i}},calculateScale:function(){var t=this.el.dataset,e=t.zoomingHeight,i=t.zoomingWidth,n=this.instance.options,s=n.customSize,o=n.scaleBase;if(!s&&e&&i)return{x:i/this.rect.width,y:e/this.rect.height};if(s&&"object"===(void 0===s?"undefined":v(s)))return{x:s.width/this.rect.width,y:s.height/this.rect.height};var r=this.rect.width/2,a=this.rect.height/2,l=k(),h={x:l.x-r,y:l.y-a},c=h.x/r,u=h.y/a,d=o+Math.min(c,u);if(s&&"string"==typeof s){var f=i||this.el.naturalWidth,p=e||this.el.naturalHeight,y=parseFloat(s)*f/(100*this.rect.width),m=parseFloat(s)*p/(100*this.rect.height);if(d>y||d>m)return{x:y,y:m}}return{x:d,y:d}}};function k(){var t=document.documentElement;return{x:Math.min(t.clientWidth,window.innerWidth)/2,y:Math.min(t.clientHeight,window.innerHeight)/2}}function S(t,e,i){["mousedown","mousemove","mouseup","touchstart","touchmove","touchend"].forEach(function(n){o(t,n,e[n],i)})}return function(){function i(t){b(this,i),this.target=Object.create(O),this.overlay=Object.create(g),this.handler=Object.create(p),this.body=document.body,this.shown=!1,this.lock=!1,this.released=!0,this.lastScrollPosition=null,this.pressTimer=null,this.options=x({},f,t),this.overlay.init(this),this.handler.init(this)}return w(i,[{key:"listen",value:function(t){if("string"==typeof t)for(var i=document.querySelectorAll(t),n=i.length;n--;)this.listen(i[n]);else"IMG"===t.tagName&&(t.style.cursor=e,o(t,"click",this.handler.click),this.options.preloadImage&&r(a(t)));return this}},{key:"config",value:function(t){return t?(x(this.options,t),this.overlay.updateStyle(this.options),this):this.options}},{key:"open",value:function(t){var e=this,i=arguments.length>1&&void 0!==arguments[1]?arguments[1]:this.options.onOpen;if(!this.shown&&!this.lock){var n="string"==typeof t?document.querySelector(t):t;if("IMG"===n.tagName){if(this.options.onBeforeOpen(n),this.target.init(n,this),!this.options.preloadImage){var s=this.target.srcOriginal;null!=s&&(this.options.onImageLoading(n),r(s,this.options.onImageLoaded))}this.shown=!0,this.lock=!0,this.target.zoomIn(),this.overlay.insert(),this.overlay.fadeIn(),o(document,"scroll",this.handler.scroll),o(document,"keydown",this.handler.keydown),this.options.closeOnWindowResize&&o(window,"resize",this.handler.resizeWindow);return o(n,u,function t(){o(n,u,t,!1),e.lock=!1,e.target.upgradeSource(),e.options.enableGrab&&S(document,e.handler,!0),i(n)}),this}}}},{key:"close",value:function(){var e=this,i=arguments.length>0&&void 0!==arguments[0]?arguments[0]:this.options.onClose;if(this.shown&&!this.lock){var n=this.target.el;this.options.onBeforeClose(n),this.lock=!0,this.body.style.cursor=t,this.overlay.fadeOut(),this.target.zoomOut(),o(document,"scroll",this.handler.scroll,!1),o(document,"keydown",this.handler.keydown,!1),this.options.closeOnWindowResize&&o(window,"resize",this.handler.resizeWindow,!1);return o(n,u,function t(){o(n,u,t,!1),e.shown=!1,e.lock=!1,e.target.downgradeSource(),e.options.enableGrab&&S(document,e.handler,!1),e.target.restoreCloseStyle(),e.overlay.remove(),i(n)}),this}}},{key:"grab",value:function(t,e){var i=arguments.length>2&&void 0!==arguments[2]?arguments[2]:this.options.scaleExtra,n=arguments.length>3&&void 0!==arguments[3]?arguments[3]:this.options.onGrab;if(this.shown&&!this.lock){var s=this.target.el;this.options.onBeforeGrab(s),this.released=!1,this.target.grab(t,e,i);return o(s,u,function t(){o(s,u,t,!1),n(s)}),this}}},{key:"move",value:function(t,e){var i=arguments.length>2&&void 0!==arguments[2]?arguments[2]:this.options.scaleExtra,n=arguments.length>3&&void 0!==arguments[3]?arguments[3]:this.options.onMove;if(this.shown&&!this.lock){this.released=!1,this.body.style.cursor=s,this.target.move(t,e,i);var r=this.target.el;return o(r,u,function t(){o(r,u,t,!1),n(r)}),this}}},{key:"release",value:function(){var e=this,i=arguments.length>0&&void 0!==arguments[0]?arguments[0]:this.options.onRelease;if(this.shown&&!this.lock){var n=this.target.el;this.options.onBeforeRelease(n),this.lock=!0,this.body.style.cursor=t,this.target.restoreOpenStyle();return o(n,u,function t(){o(n,u,t,!1),e.lock=!1,e.released=!0,i(n)}),this}}}]),i}()});

inference.py

@@ -467,16 +467,15 @@ def dock_fragments(
             log.warning("Failed on", orig_complex_graph["name"], e)
             failures += 1
 
-    # Tear down DiffDock models and datasets
-    model.cpu()
+    # Teardown
     del model
     if confidence_model is not None:
-        confidence_model.cpu()
         del confidence_model
     del test_dataset
     if confidence_test_dataset is not None:
         del confidence_test_dataset
     del test_loader
+    torch.cuda.empty_cache()
 
     docking_df = pd.concat(docking_dfs, ignore_index=True)
     # Save intermediate docking results
@@ -691,6 +690,8 @@ def select_fragment_pairs(
         ligand_residue=None,
 ):
     df = df[df['confidence'] > confidence_threshold].copy()
+    if 'ligand_conf_path' in df.columns:
+        df['ligand_conf_path'] = df['ligand_conf_path'].apply(Path)
     if 'ligand_mol' not in df.columns:
         df['ligand_mol'] = df['ligand_conf_path'].apply(read_molecule)
 
@@ -1139,9 +1140,9 @@ def generate_linkers(
         #     batch_df['out_path'] = str(out_sdf)
         #     linking_dfs.append(batch_df)
 
-    # Tear down
-    ddpm.cpu()
+    # Teardown
     del ddpm
+    torch.cuda.empty_cache()
 
     if linking_dfs:
         linking_summary_df = pd.concat(linking_dfs, ignore_index=True)

utils/inference_utils.py

@@ -62,7 +62,17 @@ def compute_esm_embeddings(model, alphabet, labels, sequences):
     repr_layers = [33]
     truncation_seq_length = 1022
 
-    dataset = FastaBatchedDataset(labels, sequences)
+    # Deduplicate by label to avoid redundant ESM computation
+    unique_labels = set()
+    unique_sequences = []
+    for label, sequence in zip(labels, sequences):
+        if label not in unique_labels:
+            unique_labels.add(label)
+            unique_sequences.append(sequence)
+    unique_labels = list(unique_labels)
+
+    # Create the dataset with unique labels and sequences
+    dataset = FastaBatchedDataset(unique_labels, unique_sequences)
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
     data_loader = torch.utils.data.DataLoader(
         dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches