thnaks

app.py

@@ -1,506 +1,506 @@
-import spaces
-import logging
-import gradio as gr
-import os
-import uuid
-from datetime import datetime
-import numpy as np
-from configs.configs_base import configs as configs_base
-from configs.configs_data import data_configs
-from configs.configs_inference import inference_configs
-from runner.inference import download_infercence_cache, update_inference_configs, infer_predict, infer_detect, InferenceRunner
-from protenix.config import parse_configs, parse_sys_args
-from runner.msa_search import update_infer_json
-from protenix.web_service.prediction_visualization import plot_best_confidence_measure, PredictionLoader
-from process_data import process_data
-import json
-from typing import Dict, List
-from Bio.PDB import MMCIFParser, PDBIO
-import tempfile
-import shutil
-from Bio import PDB
-from gradio_molecule3d import Molecule3D
-
-EXAMPLE_PATH = './examples/example.json'
-example_json=[{'sequences': [{'proteinChain': {'sequence': 'MAEVIRSSAFWRSFPIFEEFDSETLCELSGIASYRKWSAGTVIFQRGDQGDYMIVVVSGRIKLSLFTPQGRELMLRQHEAGALFGEMALLDGQPRSADATAVTAAEGYVIGKKDFLALITQRPKTAEAVIRFLCAQLRDTTDRLETIALYDLNARVARFFLATLRQIHGSEMPQSANLRLTLSQTDIASILGASRPKVNRAILSLEESGAIKRADGIICCNVGRLLSIADPEEDLEHHHHHHHH', 'count': 2}}, {'dnaSequence': {'sequence': 'CTAGGTAACATTACTCGCG', 'count': 2}}, {'dnaSequence': {'sequence': 'GCGAGTAATGTTAC', 'count': 2}}, {'ligand': {'ligand': 'CCD_PCG', 'count': 2}}], 'name': '7pzb_need_search_msa'}]
-
-# Custom CSS for styling
-custom_css = """
-#logo {
-    width: 50%;
-}
-.title {
-    font-size: 32px;
-    font-weight: bold;
-    color: #4CAF50;
-    display: flex;
-    align-items: center; /* Vertically center the logo and text */
-}
-"""
-
-
-os.environ["LAYERNORM_TYPE"] = "fast_layernorm"
-os.environ["USE_DEEPSPEED_EVO_ATTTENTION"] = "False"
-# Set environment variable in the script
-#os.environ['CUTLASS_PATH'] = './cutlass'
+# import spaces
+# import logging
+# import gradio as gr
+# import os
+# import uuid
+# from datetime import datetime
+# import numpy as np
+# from configs.configs_base import configs as configs_base
+# from configs.configs_data import data_configs
+# from configs.configs_inference import inference_configs
+# from runner.inference import download_infercence_cache, update_inference_configs, infer_predict, infer_detect, InferenceRunner
+# from protenix.config import parse_configs, parse_sys_args
+# from runner.msa_search import update_infer_json
+# from protenix.web_service.prediction_visualization import plot_best_confidence_measure, PredictionLoader
+# from process_data import process_data
+# import json
+# from typing import Dict, List
+# from Bio.PDB import MMCIFParser, PDBIO
+# import tempfile
+# import shutil
+# from Bio import PDB
+# from gradio_molecule3d import Molecule3D
+
+# EXAMPLE_PATH = './examples/example.json'
+# example_json=[{'sequences': [{'proteinChain': {'sequence': 'MAEVIRSSAFWRSFPIFEEFDSETLCELSGIASYRKWSAGTVIFQRGDQGDYMIVVVSGRIKLSLFTPQGRELMLRQHEAGALFGEMALLDGQPRSADATAVTAAEGYVIGKKDFLALITQRPKTAEAVIRFLCAQLRDTTDRLETIALYDLNARVARFFLATLRQIHGSEMPQSANLRLTLSQTDIASILGASRPKVNRAILSLEESGAIKRADGIICCNVGRLLSIADPEEDLEHHHHHHHH', 'count': 2}}, {'dnaSequence': {'sequence': 'CTAGGTAACATTACTCGCG', 'count': 2}}, {'dnaSequence': {'sequence': 'GCGAGTAATGTTAC', 'count': 2}}, {'ligand': {'ligand': 'CCD_PCG', 'count': 2}}], 'name': '7pzb_need_search_msa'}]
+
+# # Custom CSS for styling
+# custom_css = """
+# #logo {
+#     width: 50%;
+# }
+# .title {
+#     font-size: 32px;
+#     font-weight: bold;
+#     color: #4CAF50;
+#     display: flex;
+#     align-items: center; /* Vertically center the logo and text */
+# }
+# """
+
+
+# os.environ["LAYERNORM_TYPE"] = "fast_layernorm"
+# os.environ["USE_DEEPSPEED_EVO_ATTTENTION"] = "False"
+# # Set environment variable in the script
+# #os.environ['CUTLASS_PATH'] = './cutlass'
+
+# # reps = [
+# #     {
+# #         "model": 0,
+# #         "chain": "",
+# #         "resname": "",
+# #         "style": "cartoon",  # Use cartoon style
+# #         "color": "whiteCarbon",
+# #         "residue_range": "",
+# #         "around": 0,
+# #         "byres": False,
+# #         "visible": True  # Ensure this representation is visible
+# #     }
+# # ]
 
 # reps = [
 #     {
 #         "model": 0,
 #         "chain": "",
 #         "resname": "",
-#         "style": "cartoon",  # Use cartoon style
+#         "style": "cartoon",
 #         "color": "whiteCarbon",
 #         "residue_range": "",
 #         "around": 0,
 #         "byres": False,
-#         "visible": True  # Ensure this representation is visible
+#         "opacity": 0.2,
+#     },
+#     {
+#         "model": 1,
+#         "chain": "",
+#         "resname": "",
+#         "style": "cartoon",
+#         "color": "cyanCarbon",
+#         "residue_range": "",
+#         "around": 0,
+#         "byres": False,
+#         "opacity": 0.8,
 #     }
 # ]
+# ##
 
-reps = [
-    {
-        "model": 0,
-        "chain": "",
-        "resname": "",
-        "style": "cartoon",
-        "color": "whiteCarbon",
-        "residue_range": "",
-        "around": 0,
-        "byres": False,
-        "opacity": 0.2,
-    },
-    {
-        "model": 1,
-        "chain": "",
-        "resname": "",
-        "style": "cartoon",
-        "color": "cyanCarbon",
-        "residue_range": "",
-        "around": 0,
-        "byres": False,
-        "opacity": 0.8,
-    }
-]
-##
-
-
-def align_pdb_files(pdb_file_1, pdb_file_2):
-    # Load the structures
-    parser = PDB.PPBuilder()
-    io = PDB.PDBIO()
-    structure_1 = PDB.PDBParser(QUIET=True).get_structure('Structure_1', pdb_file_1)
-    structure_2 = PDB.PDBParser(QUIET=True).get_structure('Structure_2', pdb_file_2)
-
-    # Superimpose the second structure onto the first
-    super_imposer = PDB.Superimposer()
-    model_1 = structure_1[0]
-    model_2 = structure_2[0]
-
-    # Extract the coordinates from the two structures
-    atoms_1 = [atom for atom in model_1.get_atoms() if atom.get_name() == "CA"]  # Use CA atoms
-    atoms_2 = [atom for atom in model_2.get_atoms() if atom.get_name() == "CA"]
-
-    # Align the structures based on the CA atoms
-    coord_1 = [atom.get_coord() for atom in atoms_1]
-    coord_2 = [atom.get_coord() for atom in atoms_2]
+
+# def align_pdb_files(pdb_file_1, pdb_file_2):
+#     # Load the structures
+#     parser = PDB.PPBuilder()
+#     io = PDB.PDBIO()
+#     structure_1 = PDB.PDBParser(QUIET=True).get_structure('Structure_1', pdb_file_1)
+#     structure_2 = PDB.PDBParser(QUIET=True).get_structure('Structure_2', pdb_file_2)
+
+#     # Superimpose the second structure onto the first
+#     super_imposer = PDB.Superimposer()
+#     model_1 = structure_1[0]
+#     model_2 = structure_2[0]
+
+#     # Extract the coordinates from the two structures
+#     atoms_1 = [atom for atom in model_1.get_atoms() if atom.get_name() == "CA"]  # Use CA atoms
+#     atoms_2 = [atom for atom in model_2.get_atoms() if atom.get_name() == "CA"]
+
+#     # Align the structures based on the CA atoms
+#     coord_1 = [atom.get_coord() for atom in atoms_1]
+#     coord_2 = [atom.get_coord() for atom in atoms_2]
     
-    super_imposer.set_atoms(atoms_1, atoms_2)
-    super_imposer.apply(model_2)  # Apply the transformation to model_2
-
-    # Save the aligned structure back to the original file
-    io.set_structure(structure_2)  # Save the aligned structure to the second file (original file)
-    io.save(pdb_file_2)
-
-# Function to convert .cif to .pdb and save as a temporary file
-def convert_cif_to_pdb(cif_path):
-    """
-    Convert a CIF file to a PDB file and save it as a temporary file.
-
-    Args:
-        cif_path (str): Path to the input CIF file.
-
-    Returns:
-        str: Path to the temporary PDB file.
-    """
-    # Initialize the MMCIF parser
-    parser = MMCIFParser()
-    structure = parser.get_structure("protein", cif_path)
-
-    # Create a temporary file for the PDB output
-    with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as temp_file:
-        temp_pdb_path = temp_file.name
-
-        # Save the structure as a PDB file
-        io = PDBIO()
-        io.set_structure(structure)
-        io.save(temp_pdb_path)
-
-    return temp_pdb_path
-
-def plot_3d(pred_loader):
-    # Get the CIF file path for the given prediction ID
-    cif_path = sorted(pred_loader.cif_paths)[0]
-
-    # Convert the CIF file to a temporary PDB file
-    temp_pdb_path = convert_cif_to_pdb(cif_path)
-
-    return temp_pdb_path, cif_path
-
-def parse_json_input(json_data: List[Dict]) -> Dict:
-    """Convert Protenix JSON format to UI-friendly structure"""
-    components = {
-        "protein_chains": [],
-        "dna_sequences": [],
-        "ligands": [],
-        "complex_name": ""
-    }
+#     super_imposer.set_atoms(atoms_1, atoms_2)
+#     super_imposer.apply(model_2)  # Apply the transformation to model_2
+
+#     # Save the aligned structure back to the original file
+#     io.set_structure(structure_2)  # Save the aligned structure to the second file (original file)
+#     io.save(pdb_file_2)
+
+# # Function to convert .cif to .pdb and save as a temporary file
+# def convert_cif_to_pdb(cif_path):
+#     """
+#     Convert a CIF file to a PDB file and save it as a temporary file.
+
+#     Args:
+#         cif_path (str): Path to the input CIF file.
+
+#     Returns:
+#         str: Path to the temporary PDB file.
+#     """
+#     # Initialize the MMCIF parser
+#     parser = MMCIFParser()
+#     structure = parser.get_structure("protein", cif_path)
+
+#     # Create a temporary file for the PDB output
+#     with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as temp_file:
+#         temp_pdb_path = temp_file.name
+
+#         # Save the structure as a PDB file
+#         io = PDBIO()
+#         io.set_structure(structure)
+#         io.save(temp_pdb_path)
+
+#     return temp_pdb_path
+
+# def plot_3d(pred_loader):
+#     # Get the CIF file path for the given prediction ID
+#     cif_path = sorted(pred_loader.cif_paths)[0]
+
+#     # Convert the CIF file to a temporary PDB file
+#     temp_pdb_path = convert_cif_to_pdb(cif_path)
+
+#     return temp_pdb_path, cif_path
+
+# def parse_json_input(json_data: List[Dict]) -> Dict:
+#     """Convert Protenix JSON format to UI-friendly structure"""
+#     components = {
+#         "protein_chains": [],
+#         "dna_sequences": [],
+#         "ligands": [],
+#         "complex_name": ""
+#     }
     
-    for entry in json_data:
-        components["complex_name"] = entry.get("name", "")
-        for seq in entry["sequences"]:
-            if "proteinChain" in seq:
-                components["protein_chains"].append({
-                    "sequence": seq["proteinChain"]["sequence"],
-                    "count": seq["proteinChain"]["count"]
-                })
-            elif "dnaSequence" in seq:
-                components["dna_sequences"].append({
-                    "sequence": seq["dnaSequence"]["sequence"],
-                    "count": seq["dnaSequence"]["count"]
-                })
-            elif "ligand" in seq:
-                components["ligands"].append({
-                    "type": seq["ligand"]["ligand"],
-                    "count": seq["ligand"]["count"]
-                })
-    return components
-
-def create_protenix_json(input_data: Dict) -> List[Dict]:
-    """Convert UI inputs to Protenix JSON format"""
-    sequences = []
+#     for entry in json_data:
+#         components["complex_name"] = entry.get("name", "")
+#         for seq in entry["sequences"]:
+#             if "proteinChain" in seq:
+#                 components["protein_chains"].append({
+#                     "sequence": seq["proteinChain"]["sequence"],
+#                     "count": seq["proteinChain"]["count"]
+#                 })
+#             elif "dnaSequence" in seq:
+#                 components["dna_sequences"].append({
+#                     "sequence": seq["dnaSequence"]["sequence"],
+#                     "count": seq["dnaSequence"]["count"]
+#                 })
+#             elif "ligand" in seq:
+#                 components["ligands"].append({
+#                     "type": seq["ligand"]["ligand"],
+#                     "count": seq["ligand"]["count"]
+#                 })
+#     return components
+
+# def create_protenix_json(input_data: Dict) -> List[Dict]:
+#     """Convert UI inputs to Protenix JSON format"""
+#     sequences = []
     
-    for pc in input_data["protein_chains"]:
-        sequences.append({
-            "proteinChain": {
-                "sequence": pc["sequence"],
-                "count": pc["count"]
-            }
-        })
+#     for pc in input_data["protein_chains"]:
+#         sequences.append({
+#             "proteinChain": {
+#                 "sequence": pc["sequence"],
+#                 "count": pc["count"]
+#             }
+#         })
     
-    for dna in input_data["dna_sequences"]:
-        sequences.append({
-            "dnaSequence": {
-                "sequence": dna["sequence"],
-                "count": dna["count"]
-            }
-        })
+#     for dna in input_data["dna_sequences"]:
+#         sequences.append({
+#             "dnaSequence": {
+#                 "sequence": dna["sequence"],
+#                 "count": dna["count"]
+#             }
+#         })
     
-    for lig in input_data["ligands"]:
-        sequences.append({
-            "ligand": {
-                "ligand": lig["type"],
-                "count": lig["count"]
-            }
-        })
+#     for lig in input_data["ligands"]:
+#         sequences.append({
+#             "ligand": {
+#                 "ligand": lig["type"],
+#                 "count": lig["count"]
+#             }
+#         })
     
-    return [{
-        "sequences": sequences,
-        "name": input_data["complex_name"]
-    }]
-
-
-#@torch.inference_mode()
-@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
-def predict_structure(input_collector: dict):
-        #first initialize runner
-        runner = InferenceRunner(configs)
-        """Handle both input types"""
-        os.makedirs("./output", exist_ok=True)
+#     return [{
+#         "sequences": sequences,
+#         "name": input_data["complex_name"]
+#     }]
+
+
+# #@torch.inference_mode()
+# @spaces.GPU(duration=120)  # Specify a duration to avoid timeout
+# def predict_structure(input_collector: dict):
+#         #first initialize runner
+#         runner = InferenceRunner(configs)
+#         """Handle both input types"""
+#         os.makedirs("./output", exist_ok=True)
         
-        # Generate random filename with timestamp
-        random_name = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
-        save_path = os.path.join("./output", f"{random_name}.json")
-
-        print(input_collector)
-
-        # Handle JSON input
-        if input_collector["json"]:
-            # Handle different input types
-            if isinstance(input_collector["json"], str):  # Example JSON case (file path)
-                input_data = json.load(open(input_collector["json"]))
-            elif hasattr(input_collector["json"], "name"):  # File upload case
-                input_data = json.load(open(input_collector["json"].name))
-            else:  # Direct JSON data case
-                input_data = input_collector["json"]
-        else:  # Manual input case
-            input_data = create_protenix_json(input_collector["data"])
-
-        with open(save_path, "w") as f:
-            json.dump(input_data, f, indent=2)
-
-        if input_data==example_json and input_collector['watermark']==True:
-            configs.saved_path = './output/example_output/'
-        else:
-            # run msa
-            json_file = update_infer_json(save_path, './output', True)
-
-            # Run prediction
-            configs.input_json_path = json_file
-            configs.watermark = input_collector['watermark']
-            configs.saved_path = os.path.join("./output/", random_name)
-            infer_predict(runner, configs)
-            #saved_path = os.path.join('./output', f"{sample_name}", f"seed_{seed}", 'predictions')
-
-        # Generate visualizations
-        pred_loader = PredictionLoader(os.path.join(configs.saved_path, 'predictions'))
-        view3d, cif_path = plot_3d(pred_loader=pred_loader)
-        if configs.watermark:
-            pred_loader = PredictionLoader(os.path.join(configs.saved_path, 'predictions_orig'))
-            view3d_orig, _ = plot_3d(pred_loader=pred_loader)
-            align_pdb_files(view3d, view3d_orig)
-            view3d = [view3d, view3d_orig]
-        plot_best_confidence_measure(os.path.join(configs.saved_path, 'predictions'))
-        confidence_img_path = os.path.join(os.path.join(configs.saved_path, 'predictions'), "best_sample_confidence.png")
-
-        return view3d, confidence_img_path, cif_path
-
-
-logger = logging.getLogger(__name__)
-LOG_FORMAT = "%(asctime)s,%(msecs)-3d %(levelname)-8s [%(filename)s:%(lineno)s %(funcName)s] %(message)s"
-logging.basicConfig(
-    format=LOG_FORMAT,
-    level=logging.INFO,
-    datefmt="%Y-%m-%d %H:%M:%S",
-    filemode="w",
-)
-configs_base["use_deepspeed_evo_attention"] = (
-    os.environ.get("USE_DEEPSPEED_EVO_ATTTENTION", False) == "False"
-)
-arg_str = "--seeds 101 --dump_dir ./output --input_json_path ./examples/example.json --model.N_cycle 10 --sample_diffusion.N_sample 5 --sample_diffusion.N_step 200 "
-configs = {**configs_base, **{"data": data_configs}, **inference_configs}
-configs = parse_configs(
-    configs=configs,
-    arg_str=arg_str,
-    fill_required_with_null=True,
-)
-configs.load_checkpoint_path='./checkpoint.pt'
-download_infercence_cache()
-configs.use_deepspeed_evo_attention=False
-
-add_watermark = gr.Checkbox(label="Add Watermark", value=True)
-add_watermark1 = gr.Checkbox(label="Add Watermark", value=True)
-
-
-with gr.Blocks(title="FoldMark", css=custom_css) as demo:
-    with gr.Row():
-        # Use a Column to align the logo and title horizontally
-            gr.Image(value="./assets/foldmark_head.png", elem_id="logo", label="Logo", height=150, show_label=False)
-
-    with gr.Tab("Structure Predictor (JSON Upload)"):
-        # First create the upload component
-        json_upload = gr.File(label="Upload JSON", file_types=[".json"])
+#         # Generate random filename with timestamp
+#         random_name = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
+#         save_path = os.path.join("./output", f"{random_name}.json")
+
+#         print(input_collector)
+
+#         # Handle JSON input
+#         if input_collector["json"]:
+#             # Handle different input types
+#             if isinstance(input_collector["json"], str):  # Example JSON case (file path)
+#                 input_data = json.load(open(input_collector["json"]))
+#             elif hasattr(input_collector["json"], "name"):  # File upload case
+#                 input_data = json.load(open(input_collector["json"].name))
+#             else:  # Direct JSON data case
+#                 input_data = input_collector["json"]
+#         else:  # Manual input case
+#             input_data = create_protenix_json(input_collector["data"])
+
+#         with open(save_path, "w") as f:
+#             json.dump(input_data, f, indent=2)
+
+#         if input_data==example_json and input_collector['watermark']==True:
+#             configs.saved_path = './output/example_output/'
+#         else:
+#             # run msa
+#             json_file = update_infer_json(save_path, './output', True)
+
+#             # Run prediction
+#             configs.input_json_path = json_file
+#             configs.watermark = input_collector['watermark']
+#             configs.saved_path = os.path.join("./output/", random_name)
+#             infer_predict(runner, configs)
+#             #saved_path = os.path.join('./output', f"{sample_name}", f"seed_{seed}", 'predictions')
+
+#         # Generate visualizations
+#         pred_loader = PredictionLoader(os.path.join(configs.saved_path, 'predictions'))
+#         view3d, cif_path = plot_3d(pred_loader=pred_loader)
+#         if configs.watermark:
+#             pred_loader = PredictionLoader(os.path.join(configs.saved_path, 'predictions_orig'))
+#             view3d_orig, _ = plot_3d(pred_loader=pred_loader)
+#             align_pdb_files(view3d, view3d_orig)
+#             view3d = [view3d, view3d_orig]
+#         plot_best_confidence_measure(os.path.join(configs.saved_path, 'predictions'))
+#         confidence_img_path = os.path.join(os.path.join(configs.saved_path, 'predictions'), "best_sample_confidence.png")
+
+#         return view3d, confidence_img_path, cif_path
+
+
+# logger = logging.getLogger(__name__)
+# LOG_FORMAT = "%(asctime)s,%(msecs)-3d %(levelname)-8s [%(filename)s:%(lineno)s %(funcName)s] %(message)s"
+# logging.basicConfig(
+#     format=LOG_FORMAT,
+#     level=logging.INFO,
+#     datefmt="%Y-%m-%d %H:%M:%S",
+#     filemode="w",
+# )
+# configs_base["use_deepspeed_evo_attention"] = (
+#     os.environ.get("USE_DEEPSPEED_EVO_ATTTENTION", False) == "False"
+# )
+# arg_str = "--seeds 101 --dump_dir ./output --input_json_path ./examples/example.json --model.N_cycle 10 --sample_diffusion.N_sample 5 --sample_diffusion.N_step 200 "
+# configs = {**configs_base, **{"data": data_configs}, **inference_configs}
+# configs = parse_configs(
+#     configs=configs,
+#     arg_str=arg_str,
+#     fill_required_with_null=True,
+# )
+# configs.load_checkpoint_path='./checkpoint.pt'
+# download_infercence_cache()
+# configs.use_deepspeed_evo_attention=False
+
+# add_watermark = gr.Checkbox(label="Add Watermark", value=True)
+# add_watermark1 = gr.Checkbox(label="Add Watermark", value=True)
+
+
+# with gr.Blocks(title="FoldMark", css=custom_css) as demo:
+#     with gr.Row():
+#         # Use a Column to align the logo and title horizontally
+#             gr.Image(value="./assets/foldmark_head.png", elem_id="logo", label="Logo", height=150, show_label=False)
+
+#     with gr.Tab("Structure Predictor (JSON Upload)"):
+#         # First create the upload component
+#         json_upload = gr.File(label="Upload JSON", file_types=[".json"])
         
-        # Then create the example component that references it
-        gr.Examples(
-            examples=[[EXAMPLE_PATH]],
-            inputs=[json_upload],
-            label="Click to use example JSON:",
-            examples_per_page=1
-        )
+#         # Then create the example component that references it
+#         gr.Examples(
+#             examples=[[EXAMPLE_PATH]],
+#             inputs=[json_upload],
+#             label="Click to use example JSON:",
+#             examples_per_page=1
+#         )
         
-        # Rest of the components
-        upload_name = gr.Textbox(label="Complex Name (optional)")
-        upload_output = gr.JSON(label="Parsed Components")
+#         # Rest of the components
+#         upload_name = gr.Textbox(label="Complex Name (optional)")
+#         upload_output = gr.JSON(label="Parsed Components")
         
-        json_upload.upload(
-            fn=lambda f: parse_json_input(json.load(open(f.name))),
-            inputs=json_upload,
-            outputs=upload_output
-        )
-
-        # Shared prediction components
-        with gr.Row():
-            add_watermark.render()
-            submit_btn = gr.Button("Predict Structure", variant="primary")
-            #structure_view = gr.HTML(label="3D Visualization")
-
-        with gr.Row():
-            view3d = Molecule3D(label="3D Visualization", reps=reps)
-        legend = gr.Markdown("""
-        **Color Legend:**
-
-        - <span style="color:grey">Unwatermarked Structure</span>
-        - <span style="color:cyan">Watermarked Structure</span>
-        """)
-        with gr.Row():
-            cif_file = gr.File(label="Download CIF File")
-        with gr.Row():
-            confidence_plot_image = gr.Image(label="Confidence Measures")
+#         json_upload.upload(
+#             fn=lambda f: parse_json_input(json.load(open(f.name))),
+#             inputs=json_upload,
+#             outputs=upload_output
+#         )
+
+#         # Shared prediction components
+#         with gr.Row():
+#             add_watermark.render()
+#             submit_btn = gr.Button("Predict Structure", variant="primary")
+#             #structure_view = gr.HTML(label="3D Visualization")
+
+#         with gr.Row():
+#             view3d = Molecule3D(label="3D Visualization", reps=reps)
+#         legend = gr.Markdown("""
+#         **Color Legend:**
+
+#         - <span style="color:grey">Unwatermarked Structure</span>
+#         - <span style="color:cyan">Watermarked Structure</span>
+#         """)
+#         with gr.Row():
+#             cif_file = gr.File(label="Download CIF File")
+#         with gr.Row():
+#             confidence_plot_image = gr.Image(label="Confidence Measures")
         
-        input_collector = gr.JSON(visible=False)
-
-        # Map inputs to a dictionary
-        submit_btn.click(
-            fn=lambda j, w: {"json": j, "watermark": w},
-            inputs=[json_upload, add_watermark],
-            outputs=input_collector
-        ).then(
-            fn=predict_structure,
-            inputs=input_collector,
-            outputs=[view3d, confidence_plot_image, cif_file]
-        )
-
-        gr.Markdown(""" 
-        The example of the uploaded json file for structure prediction.
-        <pre>
-            [{
-        "sequences": [
-            {
-                "proteinChain": {
-                    "sequence": "MAEVIRSSAFWRSFPIFEEFDSETLCELSGIASYRKWSAGTVIFQRGDQGDYMIVVVSGRIKLSLFTPQGRELMLRQHEAGALFGEMALLDGQPRSADATAVTAAEGYVIGKKDFLALITQRPKTAEAVIRFLCAQLRDTTDRLETIALYDLNARVARFFLATLRQIHGSEMPQSANLRLTLSQTDIASILGASRPKVNRAILSLEESGAIKRADGIICCNVGRLLSIADPEEDLEHHHHHHHH",
-                    "count": 2
-                }
-            },
-            {
-                "dnaSequence": {
-                    "sequence": "CTAGGTAACATTACTCGCG",
-                    "count": 2
-                }
-            },
-            {
-                "dnaSequence": {
-                    "sequence": "GCGAGTAATGTTAC",
-                    "count": 2
-                }
-            },
-            {
-                "ligand": {
-                    "ligand": "CCD_PCG",
-                    "count": 2
-                }
-            }
-        ],
-        "name": "7pzb"
-        }]
-        </pre>
-        """)
+#         input_collector = gr.JSON(visible=False)
+
+#         # Map inputs to a dictionary
+#         submit_btn.click(
+#             fn=lambda j, w: {"json": j, "watermark": w},
+#             inputs=[json_upload, add_watermark],
+#             outputs=input_collector
+#         ).then(
+#             fn=predict_structure,
+#             inputs=input_collector,
+#             outputs=[view3d, confidence_plot_image, cif_file]
+#         )
+
+#         gr.Markdown(""" 
+#         The example of the uploaded json file for structure prediction.
+#         <pre>
+#             [{
+#         "sequences": [
+#             {
+#                 "proteinChain": {
+#                     "sequence": "MAEVIRSSAFWRSFPIFEEFDSETLCELSGIASYRKWSAGTVIFQRGDQGDYMIVVVSGRIKLSLFTPQGRELMLRQHEAGALFGEMALLDGQPRSADATAVTAAEGYVIGKKDFLALITQRPKTAEAVIRFLCAQLRDTTDRLETIALYDLNARVARFFLATLRQIHGSEMPQSANLRLTLSQTDIASILGASRPKVNRAILSLEESGAIKRADGIICCNVGRLLSIADPEEDLEHHHHHHHH",
+#                     "count": 2
+#                 }
+#             },
+#             {
+#                 "dnaSequence": {
+#                     "sequence": "CTAGGTAACATTACTCGCG",
+#                     "count": 2
+#                 }
+#             },
+#             {
+#                 "dnaSequence": {
+#                     "sequence": "GCGAGTAATGTTAC",
+#                     "count": 2
+#                 }
+#             },
+#             {
+#                 "ligand": {
+#                     "ligand": "CCD_PCG",
+#                     "count": 2
+#                 }
+#             }
+#         ],
+#         "name": "7pzb"
+#         }]
+#         </pre>
+#         """)
     
-    with gr.Tab("Structure Predictor (Manual Input)"):
-        with gr.Row():
-            complex_name = gr.Textbox(label="Complex Name")
+#     with gr.Tab("Structure Predictor (Manual Input)"):
+#         with gr.Row():
+#             complex_name = gr.Textbox(label="Complex Name")
             
-        # Replace gr.Group with gr.Accordion
-        with gr.Accordion(label="Protein Chains", open=True):
-            protein_chains = gr.Dataframe(
-                headers=["Sequence", "Count"],
-                datatype=["str", "number"],
-                row_count=1,
-                col_count=(2, "fixed")
-            )
+#         # Replace gr.Group with gr.Accordion
+#         with gr.Accordion(label="Protein Chains", open=True):
+#             protein_chains = gr.Dataframe(
+#                 headers=["Sequence", "Count"],
+#                 datatype=["str", "number"],
+#                 row_count=1,
+#                 col_count=(2, "fixed")
+#             )
             
-        # Repeat for other groups
-        with gr.Accordion(label="DNA Sequences", open=True):
-            dna_sequences = gr.Dataframe(
-                headers=["Sequence", "Count"],
-                datatype=["str", "number"],
-                row_count=1
-            )
+#         # Repeat for other groups
+#         with gr.Accordion(label="DNA Sequences", open=True):
+#             dna_sequences = gr.Dataframe(
+#                 headers=["Sequence", "Count"],
+#                 datatype=["str", "number"],
+#                 row_count=1
+#             )
             
-        with gr.Accordion(label="Ligands", open=True):
-            ligands = gr.Dataframe(
-                headers=["Ligand Type", "Count"],
-                datatype=["str", "number"],
-                row_count=1
-            )
+#         with gr.Accordion(label="Ligands", open=True):
+#             ligands = gr.Dataframe(
+#                 headers=["Ligand Type", "Count"],
+#                 datatype=["str", "number"],
+#                 row_count=1
+#             )
             
-        manual_output = gr.JSON(label="Generated JSON")
+#         manual_output = gr.JSON(label="Generated JSON")
         
-        complex_name.change(
-            fn=lambda x: {"complex_name": x},
-            inputs=complex_name,
-            outputs=manual_output
-        )
-
-        # Shared prediction components
-        with gr.Row():
-            add_watermark1.render()
-            submit_btn = gr.Button("Predict Structure", variant="primary")
-            #structure_view = gr.HTML(label="3D Visualization")
-
-        with gr.Row():
-            view3d = Molecule3D(label="3D Visualization (Gray: Unwatermarked; Cyan: Watermarked)", reps=reps)
-
-        with gr.Row():
-            cif_file = gr.File(label="Download CIF File")
-        with gr.Row():
-            confidence_plot_image = gr.Image(label="Confidence Measures")
+#         complex_name.change(
+#             fn=lambda x: {"complex_name": x},
+#             inputs=complex_name,
+#             outputs=manual_output
+#         )
+
+#         # Shared prediction components
+#         with gr.Row():
+#             add_watermark1.render()
+#             submit_btn = gr.Button("Predict Structure", variant="primary")
+#             #structure_view = gr.HTML(label="3D Visualization")
+
+#         with gr.Row():
+#             view3d = Molecule3D(label="3D Visualization (Gray: Unwatermarked; Cyan: Watermarked)", reps=reps)
+
+#         with gr.Row():
+#             cif_file = gr.File(label="Download CIF File")
+#         with gr.Row():
+#             confidence_plot_image = gr.Image(label="Confidence Measures")
         
-        input_collector = gr.JSON(visible=False)
-
-        # Map inputs to a dictionary
-        submit_btn.click(
-            fn=lambda c, p, d, l, w: {"data": {"complex_name": c, "protein_chains": p, "dna_sequences": d, "ligands": l}, "watermark": w},
-            inputs=[complex_name, protein_chains, dna_sequences, ligands, add_watermark1],
-            outputs=input_collector
-        ).then(
-            fn=predict_structure,
-            inputs=input_collector,
-            outputs=[view3d, confidence_plot_image, cif_file]
-        )
-
-    @spaces.GPU(duration=120)
-    def is_watermarked(file):
-        #first initialize runner
-        runner = InferenceRunner(configs)
-        # Generate a unique subdirectory and filename
-        unique_id = str(uuid.uuid4().hex[:8])
-        subdir = os.path.join('./output', unique_id)
-        os.makedirs(subdir, exist_ok=True)
-        filename = f"{unique_id}.cif"
-        file_path = os.path.join(subdir, filename)
+#         input_collector = gr.JSON(visible=False)
+
+#         # Map inputs to a dictionary
+#         submit_btn.click(
+#             fn=lambda c, p, d, l, w: {"data": {"complex_name": c, "protein_chains": p, "dna_sequences": d, "ligands": l}, "watermark": w},
+#             inputs=[complex_name, protein_chains, dna_sequences, ligands, add_watermark1],
+#             outputs=input_collector
+#         ).then(
+#             fn=predict_structure,
+#             inputs=input_collector,
+#             outputs=[view3d, confidence_plot_image, cif_file]
+#         )
+
+#     @spaces.GPU(duration=120)
+#     def is_watermarked(file):
+#         #first initialize runner
+#         runner = InferenceRunner(configs)
+#         # Generate a unique subdirectory and filename
+#         unique_id = str(uuid.uuid4().hex[:8])
+#         subdir = os.path.join('./output', unique_id)
+#         os.makedirs(subdir, exist_ok=True)
+#         filename = f"{unique_id}.cif"
+#         file_path = os.path.join(subdir, filename)
         
-        # Save the uploaded file to the new location
-        shutil.copy(file.name, file_path)
+#         # Save the uploaded file to the new location
+#         shutil.copy(file.name, file_path)
         
-        # Call your processing functions
-        configs.process_success = process_data(subdir)
-        configs.subdir = subdir
-        result = infer_detect(runner, configs)
-        # This function should return 'Watermarked' or 'Not Watermarked'
-        temp_pdb_path = convert_cif_to_pdb(file_path)
-        if result==False:  
-            return "Not Watermarked", temp_pdb_path
-        else:
-            return "Watermarked", temp_pdb_path
+#         # Call your processing functions
+#         configs.process_success = process_data(subdir)
+#         configs.subdir = subdir
+#         result = infer_detect(runner, configs)
+#         # This function should return 'Watermarked' or 'Not Watermarked'
+#         temp_pdb_path = convert_cif_to_pdb(file_path)
+#         if result==False:  
+#             return "Not Watermarked", temp_pdb_path
+#         else:
+#             return "Watermarked", temp_pdb_path
         
     
 
-    with gr.Tab("Watermark Detector"):
-        # First create the upload component
-        cif_upload = gr.File(label="Upload .cif", file_types=["..cif"])
+#     with gr.Tab("Watermark Detector"):
+#         # First create the upload component
+#         cif_upload = gr.File(label="Upload .cif", file_types=["..cif"])
 
-        with gr.Row():
-            cif_3d_view = Molecule3D(label="3D Visualization of Input", reps=reps)
+#         with gr.Row():
+#             cif_3d_view = Molecule3D(label="3D Visualization of Input", reps=reps)
 
-        # Prediction output
-        prediction_output = gr.Textbox(label="Prediction")
+#         # Prediction output
+#         prediction_output = gr.Textbox(label="Prediction")
 
-        # Define the interaction
-        cif_upload.change(is_watermarked, inputs=cif_upload, outputs=[prediction_output, cif_3d_view])
+#         # Define the interaction
+#         cif_upload.change(is_watermarked, inputs=cif_upload, outputs=[prediction_output, cif_3d_view])
         
-        # Example files
-        example_files = [
-        "./examples/7r6r_watermarked.cif",
-        "./examples/7pzb_unwatermarked.cif"
-        ]
+#         # Example files
+#         example_files = [
+#         "./examples/7r6r_watermarked.cif",
+#         "./examples/7pzb_unwatermarked.cif"
+#         ]
 
-        gr.Examples(examples=example_files, inputs=cif_upload)
+#         gr.Examples(examples=example_files, inputs=cif_upload)
         
 
 
     
 
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+# if __name__ == "__main__":
+#     demo.launch(share=True)