RFdiffusion3 / utils /handle_events.py
gabboud's picture
replace pdb with pdb_path var
22a41b0
raw
history blame
2.79 kB
import gradio as gr
def update_batch_choices(result):
"""
Update the batch dropdown choices once RFD3 has finished generating structures. results is the output of RFD3.
Used as event triggered by the completion of RFD3 generation.
Parameters:
----------
result: list of dicts
where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
"""
if result is None:
return gr.Dropdown(choices = [])
batches = sorted(list({d["batch"] for d in result}))
return gr.update(choices=batches, visible=True)
def update_designs(batch, result):
"""
Update the design dropdown choices once RFD3 has finished generating structures. results is the output of RFD3.
Used as event triggered by the selection of a batch in the batch dropdown.
This logic of trigerring the design dropdown update when a batch is selected comes from the reasoning that different batches may have different numbers of designs.
There is no reason to think that that will be the case since the number of designs per batch is an input to RFD3, but this might still catch some bugs.
Parameters:
----------
batch: int
the batch number selected in the batch dropdown, also takes the batch dropdown object when triggered by an event function.
result: list of dicts
where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
"""
if batch is None:
return gr.update(choices=[])
designs = sorted(list({d["design"] for d in result if d["batch"] == batch}))
return gr.update(choices=designs)
def show_pdb(batch, design, result):
"""
Show the pdb content of the selected batch and design.
Used as event triggered by the show pdb button, when a batch and design are selected in the dropdowns.
Parameters:
----------
batch: int
the batch number selected in the batch dropdown, also takes the batch dropdown object when triggered by an event function.
design: int
the design number selected in the design dropdown, also takes the design dropdown object when triggered by an event function.
result: list of dicts
where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
"""
if batch is None or design is None:
return gr.update()
pdb_path= next(d["pdb_path"] for d in result if d["batch"] == int(batch) and d["design"] == int(design))
with open(pdb_path, 'r') as f:
pdb_str = f.read()
return gr.update(value=f"Selected Batch: {batch}, Design: {design}, saved at {pdb_str}:\n {pdb_str}", visible=True)