Update rfdiffusion.toml

Author: lyskov

assets/pixi/rfdiffusion.toml

@@ -2,29 +2,57 @@
 name = "rfdiffusion"
 version = "0.0.1"
 description = "Pixi environment for https://github.com/RosettaCommons/rc RFdiffusion app"
-# platforms = ["linux-64"]
-platforms = ["linux-64", "osx-64", "osx-arm64", "win-64"]
-channels = ["conda-forge", "pytorch", "dglteam", "nvidia"]
+platforms = ["linux-64"]
+#platforms = ["linux-64", "osx-64", "osx-arm64", "win-64"]
+
+channels = ["pytorch", "nvidia", "conda-forge", "dglteam"]
 
 [dependencies]
 python = "==3.9"
 pip = "*"
 
+numpy = "1.*"
+
+pytorch = { version = "==1.12.1", channel = "pytorch" }
+pytorch-cuda = { version = "==11.6", channel = "pytorch" }
+"dgl-cuda11.6" = "*"
+
+mkl = "<2024.1"
+
+[pypi-dependencies]
+e3nn = "==0.3.3"
+wandb = "==0.12.0"
+pynvml = "==11.0.0"
+dllogger = { git = "https://github.com/NVIDIA/dllogger" }
+decorator = "==5.1.0"
+hydra-core = "==1.3.2"
+pyrsistent = "==0.19.3"
+
 [tasks]
+setup = { depends-on = ["install", "download_weights"] }
+
 install = """
 ( rm -rf rfdiffusion-repo-clone ; git clone https://github.com/RosettaCommons/RFdiffusion.git rfdiffusion-repo-clone ) \
-&& pip install \
-	dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html \
-	torch==1.12.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 \
-	e3nn==0.3.3 \
-	wandb==0.12.0 \
-	pynvml==11.0.0 \
-	git+https://github.com/NVIDIA/dllogger#egg=dllogger \
-	decorator==5.1.0 \
-	hydra-core==1.3.2 \
-	pyrsistent==0.19.3 \
 && pip install --no-cache-dir rfdiffusion-repo-clone/env/SE3Transformer \
 && pip install --no-cache-dir -e rfdiffusion-repo-clone/ --no-deps \
 """
 
-setup = { depends-on = ["install"] }
+[tasks.download_weights]
+depends-on = ["install"]
+cmd = """
+( rm -rf rfdiffusion-repo-clone/models ; mkdir -p rfdiffusion-repo-clone/models ) \
+&& wget -O rfdiffusion-repo-clone/models/Base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/Complex_base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/Complex_Fold_base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/InpaintSeq_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/InpaintSeq_Fold_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/ActiveSite_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/Base_epoch8_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/Complex_beta_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/RF_structure_prediction_weights.pt http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt \
+"""
+
+[tasks.execute]
+args = ["args"]
+cmd = "python scripts/run_inference.py {{args}}"
+cwd = "rfdiffusion-repo-clone"

src/executor/native.rs

@@ -37,16 +37,21 @@ impl Executor {
                 Ok(())
             },
         )?;
+        let new_args = spec
+            .args
+            .into_iter()
+            .map(|arg| shell_escape::escape(arg.into()).to_string())
+            .collect::<Vec<_>>()
+            .join(" ");
 
         Command::new("pixi")
             .cd(&pixi_evn_root)
             .arg("run")
             .arg("execute")
-            // .arg(shell_escape::escape(spec.args.join(" ").into()))
-            .args(spec.args)
+            .arg(new_args)
             .live()
             .exec()?;
-
+        
         Ok(())
     }