Update app.py

app.py

@@ -38,77 +38,80 @@ if not os.path.exists(model):
 
 @spaces.GPU(duration=120)
 def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples):
-    jobname = re.sub(r'[<>:"/\\|?*]', '_', jobname)
-    if jobname == "":
-        raise gr.Error("Job name empty or only invalid characters. Choose a plaintext name.")
-    os.makedirs(jobname, exist_ok=True)
-    """format Gradio Component:
-    # {"chains": [
-    #     {
-    #         "class": "DNA",
-    #         "sequence": "ATGCGT",
-    #         "chain": "A"
-    #     }
-    # ], "covMods":[]
-    # }
-    """
-    #sequences_for_msa = []
-    output = {
-    "sequences": []
-    }
-    representations = []
-    for chain in inputs["chains"]:
-        entity_type = chain["class"].lower()
-        sequence_data = {
-            entity_type: {
-                "id": chain["chain"],
-            }
+    try:
+        jobname = re.sub(r'[<>:"/\\|?*]', '_', jobname)
+        if jobname == "":
+            raise gr.Error("Job name empty or only invalid characters. Choose a plaintext name.")
+        os.makedirs(jobname, exist_ok=True)
+        """format Gradio Component:
+        # {"chains": [
+        #     {
+        #         "class": "DNA",
+        #         "sequence": "ATGCGT",
+        #         "chain": "A"
+        #     }
+        # ], "covMods":[]
+        # }
+        """
+        #sequences_for_msa = []
+        output = {
+        "sequences": []
         }
-        if entity_type in ["protein", "dna", "rna"]:
-            sequence_data[entity_type]["sequence"] = chain["sequence"]
-            if entity_type == "protein":
-                #sequences_for_msa.append(chain["sequence"])
-                if chain["msa"] == False:
-                    sequence_data[entity_type]["msa"] = f"empty"
-            representations.append({"model":0, "chain":chain["chain"], "style":"cartoon"})
-        if entity_type == "ligand":
-            if "sdf" in chain.keys():
-                if chain["sdf"]!="" and chain["name"]=="":
-                    raise gr.Error("Sorry, no SDF support yet.")
-            if "name" in chain.keys() and len(chain["name"])==3:
-                 sequence_data[entity_type]["ccd"] = chain["name"]
-            elif "smiles" in chain.keys():
-                 sequence_data[entity_type]["smiles"] = chain["smiles"]
-            else:
-                raise gr.Error("No ligand found, or not in the right format. CCD codes have 3 letters")
-                
-                
-            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. ")
-        output["sequences"].append(sequence_data)
-
-    # Convert the output to YAML
-    yaml_file_path = f"{jobname}/{jobname}.yaml"
-
-    # Write the YAML output to the file
-    with open(yaml_file_path, "w") as file:
-        yaml.dump(output, file, sort_keys=False, default_flow_style=False)
-
-    os.system(f"cat {yaml_file_path}")
-    #a3m_lines_mmseqs2 = run_mmseqs2(
-    #                sequences_for_msa,
-    #                f"./{jobname}",
-    #                use_templates=False,
-    #            )
-    #with open(f"{jobname}/msa.a3m", "w+") as fp:
-    #    fp.writelines(a3m_lines_mmseqs2)
-
-    os.system(f"boltz predict {jobname}/{jobname}.yaml --use_msa_server --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 Molecule3D(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb", label="Output", reps=representations)
+        representations = []
+        for chain in inputs["chains"]:
+            entity_type = chain["class"].lower()
+            sequence_data = {
+                entity_type: {
+                    "id": chain["chain"],
+                }
+            }
+            if entity_type in ["protein", "dna", "rna"]:
+                sequence_data[entity_type]["sequence"] = chain["sequence"]
+                if entity_type == "protein":
+                    #sequences_for_msa.append(chain["sequence"])
+                    if chain["msa"] == False:
+                        sequence_data[entity_type]["msa"] = f"empty"
+                representations.append({"model":0, "chain":chain["chain"], "style":"cartoon"})
+            if entity_type == "ligand":
+                if "sdf" in chain.keys():
+                    if chain["sdf"]!="" and chain["name"]=="":
+                        raise gr.Error("Sorry, no SDF support yet.")
+                if "name" in chain.keys() and len(chain["name"])==3:
+                     sequence_data[entity_type]["ccd"] = chain["name"]
+                elif "smiles" in chain.keys():
+                     sequence_data[entity_type]["smiles"] = chain["smiles"]
+                else:
+                    raise gr.Error("No ligand found, or not in the right format. CCD codes have 3 letters")
+                    
+                    
+                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. ")
+            output["sequences"].append(sequence_data)
+    
+        # Convert the output to YAML
+        yaml_file_path = f"{jobname}/{jobname}.yaml"
+    
+        # Write the YAML output to the file
+        with open(yaml_file_path, "w") as file:
+            yaml.dump(output, file, sort_keys=False, default_flow_style=False)
+    
+        os.system(f"cat {yaml_file_path}")
+        #a3m_lines_mmseqs2 = run_mmseqs2(
+        #                sequences_for_msa,
+        #                f"./{jobname}",
+        #                use_templates=False,
+        #            )
+        #with open(f"{jobname}/msa.a3m", "w+") as fp:
+        #    fp.writelines(a3m_lines_mmseqs2)
+    
+        os.system(f"boltz predict {jobname}/{jobname}.yaml --use_msa_server --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 Molecule3D(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb", label="Output", reps=representations)
+    except Exception as e:
+        raise gr.Error(f"failed with error:{e}")
 
 with gr.Blocks() as blocks:
     gr.Markdown("# Boltz-1")