Upload app.py with huggingface_hub

app.py

@@ -1,383 +1,64 @@
-import spaces
 import gradio as gr
-
 import torch
-import numpy as np
-import pandas as pd
-import random
-import io
-import imageio
-import os
-import tempfile
-import atexit
-import glob
-import csv
-from datetime import datetime
-import json
-
 from rdkit import Chem
 from rdkit.Chem import Draw
-
-from evaluator import Evaluator
-# from loader import load_graph_decoder
-
-### load model start
 from graph_decoder.diffusion_model import GraphDiT
-def count_parameters(model):
-    r"""
-    Returns the number of trainable parameters and number of all parameters in the model.
-    """
-    trainable_params, all_param = 0, 0
-    for param in model.parameters():
-        num_params = param.numel()
-        all_param += num_params
-        if param.requires_grad:
-            trainable_params += num_params
-
-    return trainable_params, all_param
 
+# Load the model
 def load_graph_decoder(path='model_labeled'):
-    model_config_path = f"{path}/config.yaml"
-    data_info_path = f"{path}/data.meta.json"
-
     model = GraphDiT(
-        model_config_path=model_config_path,
-        data_info_path=data_info_path,
+        model_config_path=f"{path}/config.yaml",
+        data_info_path=f"{path}/data.meta.json",
         model_dtype=torch.float32,
     )
     model.init_model(path)
     model.disable_grads()
-
-    trainable_params, all_param = count_parameters(model)
-    param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
-        path, trainable_params, all_param, 100 * trainable_params / all_param
-    )
-    print(param_stats)
     return model
-### load model end
-
-# Load the CSV data
-known_labels = pd.read_csv('data/known_labels.csv')
-knwon_smiles = pd.read_csv('data/known_polymers.csv')
-
-all_properties = ['CH4', 'CO2', 'H2', 'N2', 'O2']
-
-# Initialize evaluators
-evaluators = {prop: Evaluator(f'evaluators/{prop}.joblib', prop) for prop in all_properties}
-
-# Get min and max values for each property
-property_ranges = {prop: (known_labels[prop].min(), known_labels[prop].max()) for prop in all_properties}
-
-# Create a temporary directory for GIFs
-temp_dir = tempfile.mkdtemp(prefix="polymer_gifs_")
-
-def cleanup_temp_files():
-    """Clean up temporary GIF files on exit."""
-    for file in glob.glob(os.path.join(temp_dir, "*.gif")):
-        try:
-            os.remove(file)
-        except Exception as e:
-            print(f"Error deleting {file}: {e}")
-    try:
-        os.rmdir(temp_dir)
-    except Exception as e:
-        print(f"Error deleting temporary directory {temp_dir}: {e}")
 
-# Register the cleanup function to be called on exit
-atexit.register(cleanup_temp_files)
+model = load_graph_decoder()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-def random_properties():
-    return known_labels[all_properties].sample(1).values.tolist()[0]
-
-def load_model(model_choice):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = load_graph_decoder(path=model_choice)
-    return (model, device)
-
-# Create a flagged folder if it doesn't exist
-flagged_folder = "flagged"
-os.makedirs(flagged_folder, exist_ok=True)
-
-def save_interesting_log(smiles, properties, suggested_properties):
-    """Save interesting polymer data to a CSV file."""
-    log_file = os.path.join(flagged_folder, "log.csv")
-    file_exists = os.path.isfile(log_file)
-    
-    with open(log_file, 'a', newline='') as csvfile:
-        fieldnames = ['timestamp', 'smiles'] + all_properties + [f'suggested_{prop}' for prop in all_properties]
-        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-        
-        if not file_exists:
-            writer.writeheader()
-        
-        log_data = {
-            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            'smiles': smiles,
-            **{prop: value for prop, value in zip(all_properties, properties)},
-            **{f'suggested_{prop}': value for prop, value in suggested_properties.items()}
-        }
-        writer.writerow(log_data)
-
-@spaces.GPU(duration=75)
-def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
-    print('in generate_graph')
-    model, device = model_state
-    
+def generate_polymer(CH4, CO2, H2, N2, O2, guidance_scale):
     properties = [CH4, CO2, H2, N2, O2]
     
-    def is_nan_like(x):
-        return x == 0 or x == '' or (isinstance(x, float) and np.isnan(x))
-    
-    properties = [None if is_nan_like(prop) else prop for prop in properties]
-    
-    nan_message = "The following gas properties were treated as NaN: "
-    nan_gases = [gas for gas, prop in zip(all_properties, properties) if prop is None]
-    nan_message += ", ".join(nan_gases) if nan_gases else "None"
-
-    num_nodes = None if num_nodes == 0 else num_nodes
-    
-    for _ in range(repeating_time):
-        # try:
+    try:
         model.to(device)
-        generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
-
-        # Create GIF if img_list is available
-        gif_path = None
-        if img_list and len(img_list) > 0:
-            imgs = [np.array(pil_img) for pil_img in img_list]
-            imgs.extend([imgs[-1]] * 10)
-            gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
-            imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
+        generated_molecule, _ = model.generate(properties, device=device, guide_scale=guidance_scale)
         
         if generated_molecule is not None:
             mol = Chem.MolFromSmiles(generated_molecule)
             if mol is not None:
                 standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-                is_novel = standardized_smiles not in knwon_smiles['SMILES'].values
-                novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
                 img = Draw.MolToImage(mol)
-                
-                # Evaluate the generated molecule
-                suggested_properties = {}
-                for prop, evaluator in evaluators.items():
-                    suggested_properties[prop] = evaluator([standardized_smiles])[0]
-                
-                suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
-                
-                return (
-                    f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
-                    f"**{nan_message}**\n\n"
-                    f"**{novelty_status}**\n\n"
-                    f"**Suggested Properties:**\n{suggested_properties_text}",
-                    img,
-                    gif_path,
-                    properties,  # Add this
-                    suggested_properties  # Add this
-                )
-        else:
-            return (
-                f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
-                None,
-                gif_path,
-                properties,
-                None,
-            )
-        # except Exception as e:
-        #     print(f"Error in generation: {e}")
-        #     continue
+                return standardized_smiles, img
+    except Exception as e:
+        print(f"Error in generation: {e}")
     
-    return f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**", None, None
-
-def set_random_properties():
-    return random_properties()
-
-# Create a mapping of internal names to display names
-model_name_mapping = {
-    "model_all": "Graph DiT (trained on labeled + unlabeled)",
-    "model_labeled": "Graph DiT (trained on labeled)"
-}
-
-def numpy_to_python(obj):
-    if isinstance(obj, np.integer):
-        return int(obj)
-    elif isinstance(obj, np.floating):
-        return float(obj)
-    elif isinstance(obj, np.ndarray):
-        return obj.tolist()
-    elif isinstance(obj, list):
-        return [numpy_to_python(item) for item in obj]
-    elif isinstance(obj, dict):
-        return {k: numpy_to_python(v) for k, v in obj.items()}
-    else:
-        return obj
-        
-def on_generate(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
-    result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps)
-    # Check if the generation was successful
-    if result[0].startswith("**Generated polymer SMILES:**"):
-        smiles = result[0].split("**Generated polymer SMILES:** `")[1].split("`")[0]
-        properties = json.dumps(numpy_to_python(result[3]))
-        suggested_properties = json.dumps(numpy_to_python(result[4]))
-        # Return the result with an enabled feedback button
-        return [*result[:3], smiles, properties, suggested_properties, gr.Button(interactive=True)]
-    else:
-        # Return the result with a disabled feedback button
-        return [*result[:3], "", "[]", "[]", gr.Button(interactive=False)]
-
-def process_feedback(checkbox_value, smiles, properties, suggested_properties):
-    if checkbox_value:
-        # Check if properties and suggested_properties are already Python objects
-        if isinstance(properties, str):
-            properties = json.loads(properties)
-        if isinstance(suggested_properties, str):
-            suggested_properties = json.loads(suggested_properties)
-        
-        save_interesting_log(smiles, properties, suggested_properties)
-        return gr.Textbox(value="Thank you for your feedback! This polymer has been saved to our interesting polymers log.", visible=True)
-    else:
-        return gr.Textbox(value="Thank you for your feedback!", visible=True)
-
-# ADD THIS FUNCTION
-def reset_feedback_button():
-    return gr.Button(interactive=False)
-
-# Create the Gradio interface using Blocks
-with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
-    # Navigation Bar
-    with gr.Row(elem_id="navbar"):
-        gr.Markdown("""
-        <div style="text-align: center;">
-            <h1>🔗🔬 Polymer Design with GraphDiT</h1>
-            <div style="display: flex; gap: 20px; justify-content: center; align-items: center; margin-top: 10px;">
-                <a href="https://github.com/liugangcode/Graph-DiT" target="_blank" style="display: flex; align-items: center; gap: 5px; text-decoration: none; color: inherit;">
-                    <img src="https://img.icons8.com/ios-glyphs/30/000000/github.png" alt="GitHub" />
-                    <span>View Code</span>
-                </a>
-                <a href="https://arxiv.org/abs/2401.13858" target="_blank" style="text-decoration: none; color: inherit;">
-                    📄 View Paper
-                </a>
-            </div>
-        </div>
-        """)
-
-    # Main Description
-    gr.Markdown("""
-    ## Introduction
-
-    Input the desired gas barrier properties for CH₄, CO₂, H₂, N₂, and O₂ to generate novel polymer structures. The results are visualized as molecular graphs and represented by SMILES strings if they are successfully generated. Note: Gas barrier values set to 0 will be treated as `NaN` (unconditionally). If the generation fails, please retry or increase the number of repetition attempts.
-    """)
-
-    # Model Selection
-    model_choice = gr.Radio(
-        choices=list(model_name_mapping.values()),
-        label="Model Zoo",
-        # value="Graph DiT (trained on labeled + unlabeled)"
-        value="Graph DiT (trained on labeled)"
-    )
-
-    # Model Description Accordion
-    with gr.Accordion("🔍 Model Description", open=False):
-        gr.Markdown("""
-        ### GraphDiT: Graph Diffusion Transformer
-
-        GraphDiT is a graph diffusion model designed for targeted molecular generation. It employs a conditional diffusion process to iteratively refine molecular structures based on user-specified properties.
-
-        We have collected a labeled polymer database for gas permeability from [Membrane Database](https://research.csiro.au/virtualscreening/membrane-database-polymer-gas-separation-membranes/). Additionally, we utilize unlabeled polymer structures from [PolyInfo](https://polymer.nims.go.jp/).
-
-        The gas permeability ranges from 0 to over ten thousand, with only hundreds of labeled data points, making this task particularly challenging.
-
-        We are actively working on improving the model. We welcome any feedback regarding model usage or suggestions for improvement.
-
-        #### Currently, we have two variants of Graph DiT:
-        - **Graph DiT (trained on labeled + unlabeled)**: This model uses both labeled and unlabeled data for training, potentially leading to more diverse/novel polymer generation.
-        - **Graph DiT (trained on labeled)**: This model is trained exclusively on labeled data, which may result in higher validity but potentially less diverse/novel outputs.
-        """)
-
-    # Citation Accordion
-    with gr.Accordion("📄 Citation", open=False):
-        gr.Markdown("""
-        If you use this model or interface useful, please cite the following paper:
-        ```bibtex
-        @article{graphdit2024,
-          title={Graph Diffusion Transformers for Multi-Conditional Molecular Generation},
-          author={Liu, Gang and Xu, Jiaxin and Luo, Tengfei and Jiang, Meng},
-          journal={NeurIPS},
-          year={2024},
-        }
-        ```
-        """)
-
-    model_state = gr.State(lambda: load_model("model_labeled"))
-
-    with gr.Row():
-        CH4_input = gr.Slider(0, property_ranges['CH4'][1], value=2.5, label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]")
-        CO2_input = gr.Slider(0, property_ranges['CO2'][1], value=15.4, label=f"CO₂ (Barrier) [0-{property_ranges['CO2'][1]:.1f}]")
-        H2_input = gr.Slider(0, property_ranges['H2'][1], value=21.0, label=f"H₂ (Barrier) [0-{property_ranges['H2'][1]:.1f}]")
-        N2_input = gr.Slider(0, property_ranges['N2'][1], value=1.5, label=f"N₂ (Barrier) [0-{property_ranges['N2'][1]:.1f}]")
-        O2_input = gr.Slider(0, property_ranges['O2'][1], value=2.8, label=f"O₂ (Barrier) [0-{property_ranges['O2'][1]:.1f}]")
+    return "Generation failed", None
 
+# Create the Gradio interface
+with gr.Blocks(title="Simplified Polymer Design") as iface:
+    gr.Markdown("## Polymer Design with GraphDiT")
+    
     with gr.Row():
-        guidance_scale = gr.Slider(1, 3, value=2, label="Guidance Scale from Properties")
-        num_nodes = gr.Slider(0, 50, step=1, value=0, label="Number of Nodes (0 for Random, Larger Graphs Take More Time)")
-        repeating_time = gr.Slider(1, 10, step=1, value=3, label="Repetition Until Success")
-        num_chain_steps = gr.Slider(0, 499, step=1, value=50, label="Number of Diffusion Steps to Visualize (Larger Numbers Take More Time)")
-        fps = gr.Slider(0.25, 10, step=0.25, value=5, label="Frames Per Second")
+        CH4_input = gr.Slider(0, 100, value=2.5, label="CH₄ (Barrier)")
+        CO2_input = gr.Slider(0, 100, value=15.4, label="CO₂ (Barrier)")
+        H2_input = gr.Slider(0, 100, value=21.0, label="H₂ (Barrier)")
+        N2_input = gr.Slider(0, 100, value=1.5, label="N₂ (Barrier)")
+        O2_input = gr.Slider(0, 100, value=2.8, label="O₂ (Barrier)")
+        guidance_scale = gr.Slider(1, 3, value=2, label="Guidance Scale")
 
-    with gr.Row():
-        random_btn = gr.Button("🔀 Randomize Properties (from Labeled Data)")
-        generate_btn = gr.Button("🚀 Generate Polymer")
+    generate_btn = gr.Button("Generate Polymer")
 
     with gr.Row():
-        result_text = gr.Textbox(label="📝 Generation Result")
-        result_image = gr.Image(label="Final Molecule Visualization", type="pil")
-        result_gif = gr.Image(label="Generation Process Visualization", type="filepath", format="gif")
-
-    with gr.Row() as feedback_row:
-        feedback_btn = gr.Button("🌟 I think this polymer is interesting!", visible=True, interactive=False)
-        feedback_result = gr.Textbox(label="Feedback Result", visible=False)
-    
-    # Add model switching functionality
-    def switch_model(choice):
-        # Convert display name back to internal name
-        internal_name = next(key for key, value in model_name_mapping.items() if value == choice)
-        return load_model(internal_name)
-
-    model_choice.change(switch_model, inputs=[model_choice], outputs=[model_state])
-
-    # Hidden components to store generation data
-    hidden_smiles = gr.Textbox(visible=False)
-    hidden_properties = gr.JSON(visible=False)
-    hidden_suggested_properties = gr.JSON(visible=False)
-    
-    # Set up event handlers
-    random_btn.click(
-        set_random_properties,
-        outputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input]
-    )
+        result_smiles = gr.Textbox(label="Generated SMILES")
+        result_image = gr.Image(label="Molecule Visualization", type="pil")
 
     generate_btn.click(
-        on_generate,
-        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps],
-        outputs=[result_text, result_image, result_gif, hidden_smiles, hidden_properties, hidden_suggested_properties, feedback_btn]
-    )
-
-    feedback_btn.click(
-        process_feedback,
-        inputs=[gr.Checkbox(value=True, visible=False), hidden_smiles, hidden_properties, hidden_suggested_properties],
-        outputs=[feedback_result]
-    ).then(
-        lambda: gr.Button(interactive=False),
-        outputs=[feedback_btn]
+        generate_polymer,
+        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale],
+        outputs=[result_smiles, result_image]
     )
-    
-    CH4_input.change(reset_feedback_button, outputs=[feedback_btn])
-    CO2_input.change(reset_feedback_button, outputs=[feedback_btn])
-    H2_input.change(reset_feedback_button, outputs=[feedback_btn])
-    N2_input.change(reset_feedback_button, outputs=[feedback_btn])
-    O2_input.change(reset_feedback_button, outputs=[feedback_btn])
-    random_btn.click(reset_feedback_button, outputs=[feedback_btn])
 
-# Launch the interface
 if __name__ == "__main__":
-    # iface.launch(share=True)
-    iface.launch(share=False)
+    iface.launch()