Update app.py

app.py

@@ -54,6 +54,7 @@ def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples)
     output = {
     "sequences": []
     }
+    representations = []
     for chain in inputs["chains"]:
         entity_type = chain["class"].lower()
         sequence_data = {
@@ -66,6 +67,7 @@ def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples)
             if entity_type == "protein":
                 sequences_for_msa.append(chain["sequence"])
                 sequence_data[entity_type]["msa"] = f"{jobname}/msa.a3m"
+            representations.append({"model":0, "chain":chain["chain"], "style":"cartoon"})
         if entity_type == "ligand":
             if "sdf" in chain.keys():
                 raise gr.Error("Sorry no SDF support yet")
@@ -73,6 +75,8 @@ def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples)
                  sequence_data[entity_type]["ccd"] = chain["name"]
             if "smiles" in chain.keys():
                  sequence_data[entity_type]["smiles"] = chain["smiles"]
+                
+            representations.append({"model":0, "chain":chain["chain"], "style":"stick", "color":"greenCarbon"})
 
         if len(inputs["covMods"])>0:
             raise gr.Error("Sorry, covMods not supported yet. Coming soon. ")
@@ -97,16 +101,16 @@ def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples)
     os.system(f"boltz predict {jobname}/{jobname}.yaml --out_dir {jobname} --recycling_steps {recycling_steps} --sampling_steps {sampling_steps} --diffusion_samples {diffusion_samples} --override --output_format pdb")
     print(os.listdir(jobname))
     print(os.listdir(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/"))
-    return f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb"
+    return Molecule3D(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb", label="Output", reps=representations)
 
 with gr.Blocks() as blocks:
     gr.Markdown("# Boltz-1")
     gr.Markdown("""Open GUI for running [Boltz-1 model](https://github.com/jwohlwend/boltz/) <br>
     Key components:
-    - MMSeqs2 Webserver Mirdita et al. 
-    - Boltz-1 Model Wohlwend et al.
-    - Gradio Custom Components Molecule3D/Cofolding Input Dürr S.
-    - 3dmol.js Rego & Koes 
+    - MMSeqs2 Webserver [Mirdita et al.](https://www.nature.com/articles/s41592-022-01488-1)
+    - Boltz-1 Model [Wohlwend et al.](https://github.com/jwohlwend/boltz/)
+    - Gradio Custom Components [Molecule3D](https://huggingface.co/spaces/simonduerr/gradio_molecule3d)/[Cofolding Input](https://huggingface.co/spaces/simonduerr/gradio_cofoldinginput) by myself 
+    - [3dmol.js Rego & Koes](https://academic.oup.com/bioinformatics/article/31/8/1322/213186)
     
     Note: This is an alpha: Some things like covalent modifications or using sdf files don't work yet. You can a Docker image of this on your local infrastructure easily using:
     `docker run -it -p 7860:7860 --platform=linux/amd64 --gpus all registry.hf.space/simonduerr-boltz-1:latest python app.py`