update

README.md

@@ -1,5 +1,5 @@
 ---
-title: GT4SD - Diffusers (image)
+title: GT4SD - GeoDiff
 emoji: 💡
 colorFrom: green
 colorTo: blue

app.py

@@ -1,36 +1,46 @@
 import logging
 import pathlib
+import pickle
 import gradio as gr
+from typing import Dict, Any
 import pandas as pd
 from gt4sd.algorithms.generation.diffusion import (
     DiffusersGenerationAlgorithm,
-    DDPMGenerator,
-    DDIMGenerator,
-    ScoreSdeGenerator,
-    LDMTextToImageGenerator,
-    LDMGenerator,
-    StableDiffusionGenerator,
+    GeoDiffGenerator,
 )
 from gt4sd.algorithms.registry import ApplicationsRegistry
+from utils import draw_grid_generate
+from rdkit import Chem
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
 
-def run_inference(model_type: str, prompt: str):
+def run_inference(
+    algorithm_version: str,
+    prompt_file: str,
+    prompt_id: int,
+    number_of_samples: int,
+):
 
-    if prompt == "":
-        config = eval(f"{model_type}()")
+    # Read file:
+    with open(prompt_file.name, "rb") as f:
+        prompts = pickle.load(f)
+
+    if all(isinstance(x, str) for x in prompts.keys()):
+        prompt = prompts[prompt_id]
     else:
-        config = eval(f"{model_type}(prompt={prompt})")
-    if config.modality != "token2image" and prompt != "":
-        raise ValueError(
-            f"{model_type} is an unconditional generative model, please remove prompt (not={prompt})"
-        )
+        prompt = prompts
+
+    config = GeoDiffGenerator(
+        algorithm_version=algorithm_version,
+        prompt=prompt,
+    )
     model = DiffusersGenerationAlgorithm(config)
-    image = list(model.sample(1))[0]
+    results = list(model.sample(number_of_samples))
+    smiles = [Chem.MolToSmiles(m) for m in results]
 
-    return image
+    return draw_grid_generate(samples=smiles, n_cols=5)
 
 
 if __name__ == "__main__":
@@ -38,17 +48,16 @@ if __name__ == "__main__":
     # Preparation (retrieve all available algorithms)
     all_algos = ApplicationsRegistry.list_available()
     algos = [
-        x["algorithm_application"]
-        for x in list(filter(lambda x: "Diff" in x["algorithm_name"], all_algos))
+        x["algorithm_version"]
+        for x in list(
+            filter(lambda x: "GeoDiff" in x["algorithm_application"], all_algos)
+        )
     ]
-    algos = [a for a in algos if not "GeoDiff" in a]
 
     # Load metadata
     metadata_root = pathlib.Path(__file__).parent.joinpath("model_cards")
 
-    examples = pd.read_csv(metadata_root.joinpath("examples.csv"), header=None).fillna(
-        ""
-    )
+    examples = [[algos[0], metadata_root.joinpath("mol_dct.pkl"), 2]]
 
     with open(metadata_root.joinpath("article.md"), "r") as f:
         article = f.read()
@@ -57,16 +66,18 @@ if __name__ == "__main__":
 
     demo = gr.Interface(
         fn=run_inference,
-        title="Diffusion-based image generators",
+        title="GeoDiff",
         inputs=[
             gr.Dropdown(
-                algos, label="Diffusion model", value="StableDiffusionGenerator"
+                algos, label="GeoDiff version", value="fusing/gfn-molecule-gen-drugs"
             ),
-            gr.Textbox(label="Text prompt", placeholder="A blue tree", lines=1),
+            gr.File(file_types=[".pkl"], label="GeoDiff prompt"),
+            gr.Number(value=0, label="Prompt ID", precision=0),
+            gr.Slider(minimum=1, maximum=5, value=2, label="Number of samples", step=1),
         ],
-        outputs=gr.outputs.Image(type="pil"),
+        outputs=gr.HTML(label="Output"),
         article=article,
         description=description,
-        examples=examples.values.tolist(),
+        examples=examples,
     )
     demo.launch(debug=True, show_error=True)