main.py

import logging
import os.path

import gradio as gr

import run_utils


def run_wrapper(protein_file, ligand_file, *args, **kwargs) -> str:
    return run_utils.run_cli_command(protein_file.name, ligand_file.name,
                                     *args, **kwargs)


def run():
    iface = gr.Interface(
        fn=run_wrapper,
        inputs=[
            gr.File(label="Protein PDB", file_types=[".pdb"]),
            gr.File(label="Ligand SDF", file_types=[".sdf"]),
            gr.Number(
                label="Samples Per Complex",
                value=1,
                minimum=1,
                maximum=100,
                precision=0,
            ),
            gr.Checkbox(label="Keep Local Structures", value=True),
            gr.Checkbox(label="Save Visualisation", value=True),
        ],
        outputs=gr.File(label="Result"),
    )

    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)


def set_env_variables():
    if "DiffDock-Pocket-Dir" not in os.environ:
        work_dir = os.path.abspath(os.path.join("../DiffDock-Pocket"))
        if os.path.exists(work_dir):
            os.environ["DiffDock-Pocket-Dir"] = work_dir
        else:
            raise ValueError(f"DiffDock-Pocket-Dir {work_dir} not found")

    if "LOG_LEVEL" not in os.environ:
        os.environ["LOG_LEVEL"] = "INFO"


if __name__ == "__main__":
    set_env_variables()
    run_utils.configure_logging()

    run()