Include json config

app.py

@@ -46,28 +46,29 @@ def create_fasta_file(sequence: str, name: Optional[str] = None) -> str:
     
     return file_name
 
-# Function to create a custom JSON config file
-def create_custom_config(
-    num_trunk_recycles: int = 3,
-    num_diffn_timesteps: int = 200,
-    seed: int = 42,
-    use_esm_embeddings: bool = True,
-    use_msa_server: bool = True,
-    output_file: Optional[str] = None
-) -> str:
-    """Create a custom JSON configuration file for Chai1 inference.
+# Function to create a JSON file
+def create_json_config(
+    num_diffn_timesteps: int,
+    num_trunk_recycles: int,
+    seed: int,
+    options: list
+    ) -> str:
+    """Create a JSON configuration file from the Gradio interface inputs.
     
     Args:
-        num_trunk_recycles (int, optional): Number of trunk recycles. Defaults to 3.
-        num_diffn_timesteps (int, optional): Number of diffusion timesteps. Defaults to 200.
-        seed (int, optional): Random seed. Defaults to 42.
-        use_esm_embeddings (bool, optional): Whether to use ESM embeddings. Defaults to True.
-        use_msa_server (bool, optional): Whether to use MSA server. Defaults to True.
-        output_file (str, optional): Path to save the config file. If None, saves to default location.
+        num_diffn_timesteps (int): Number of diffusion timesteps from slider
+        num_trunk_recycles (int): Number of trunk recycles from slider
+        seed (int): Random seed from slider
+        options (list): List of selected options from checkbox group
     
     Returns:
-        str: Path to the created config file
+        str: Name of the created JSON file
     """
+    # Convert checkbox options to boolean flags
+    use_esm_embeddings = "ESM_embeddings" in options
+    use_msa_server = "MSA_server" in options
+    
+    # Create config dictionary
     config = {
         "num_trunk_recycles": num_trunk_recycles,
         "num_diffn_timesteps": num_diffn_timesteps,
@@ -76,13 +77,17 @@ def create_custom_config(
         "use_msa_server": use_msa_server
     }
     
-    if output_file is None:
-        output_file = here / "inputs" / "chai1_custom_inference.json"
-        
-    with open(output_file, "w") as f:
+    # Generate a unique file name
+    unique_id = hashlib.sha256(uuid4().bytes).hexdigest()[:8]
+    file_name = f"chai1_{unique_id}_config.json"
+    file_path = here / "inputs" / file_name
+    
+    # Write the JSON file
+    with open(file_path, "w") as f:
         json.dump(config, f, indent=4)
-        
-    return str(output_file)
+    
+    return file_name
+
 
 # Function to compute Chai1 inference
 def compute_Chai1(
@@ -115,6 +120,7 @@ def compute_Chai1(
         # Define inference config file
         if not inference_config_file:
             inference_config_file = here / "inputs" / "chai1_quick_inference.json"
+        inference_config_file = here / "inputs" / inference_config_file
         print(f"🧬 loading Chai inference config from {inference_config_file}")
         inference_config = json.loads(Path(inference_config_file).read_text())
 
@@ -154,11 +160,24 @@ with gr.Blocks() as demo:
     """) 
     
     with gr.Tab("Configuration 📦"):
-        text_input = gr.Textbox(placeholder="Fasta format sequences", label="Fasta content", lines=10)
-        text_output = gr.Textbox(placeholder="Fasta file name", label="Fasta file name")
-        text_button = gr.Button("Create Fasta file")
-        text_button.click(fn=create_fasta_file, inputs=[text_input], outputs=[text_output])
         
+        with gr.Row():
+            with gr.Column(scale=1):
+                slider_nb = gr.Slider(1, 500, value=200, label="Number of diffusion time steps", info="Choose the number of diffusion time steps for the simulation", step=1, interactive=True, elem_id="num_iterations")
+                slider_trunk = gr.Slider(1, 5, value=3, label="Number of trunk recycles", info="Choose the number of iterations for the simulation", step=1, interactive=True, elem_id="trunk_number")
+                slider_seed = gr.Slider(1, 100, value=42, label="Seed", info="Choose the seed", step=1, interactive=True, elem_id="seed")
+                check_options = gr.CheckboxGroup(["ESM_embeddings", "MSA_server"], value=["ESM_embeddings",], label="Additionnal options", info="Options to use ESM embeddings and MSA server", elem_id="options")
+                json_output = gr.Textbox(placeholder="Config file name", label="Config file name")
+                button_json = gr.Button("Create Config file")
+                button_json.click(fn=create_json_config, inputs=[slider_nb, slider_trunk, slider_seed, check_options], outputs=[json_output])
+                
+            with gr.Column(scale=1):   
+                text_input = gr.Textbox(placeholder="Fasta format sequences", label="Fasta content", lines=10)
+                text_output = gr.Textbox(placeholder="Fasta file name", label="Fasta file name")
+                text_button = gr.Button("Create Fasta file")
+                text_button.click(fn=create_fasta_file, inputs=[text_input], outputs=[text_output])
+        
+
         gr.Markdown(
         """
         You can input a Fasta file containing the sequence of the molecule you want to simulate.