fix(samplers): reconcile FlowModelWrapper.step call with Protocol

src/sampleworks/cli/guidance.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import sys
-from loguru import logger
 
 from sampleworks.utils.guidance_script_arguments import GuidanceConfig
 

src/sampleworks/core/samplers/edm.py

@@ -422,10 +422,8 @@ def step(
         # t_hat will be float if check_context didn't raise
         # Use no_grad when gradients aren't needed to avoid memory overhead from
         # gradient checkpointing holding intermediate activations
-        # TODO testing adding eps to signature for use with Protpardelle-1c, if successful,
-        #   I need to modify the Protocol itself. @Michael Anzuoni
         with torch.set_grad_enabled(allow_gradients):
-            x_hat_0 = model_wrapper.step(noisy_state, t_hat, eps, features=features)
+            x_hat_0 = model_wrapper.step(noisy_state, t_hat, features=features)
 
         reconciler = (
             context.reconciler.to(torch.as_tensor(x_hat_0).device)

src/sampleworks/models/protpardelle/wrapper.py

@@ -439,7 +439,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[ProtpardelleConditi
 
     def _atom37_indices_from_atom_array(
         self, atom_array
-    ) -> tuple[Int[Tensor, "atoms"], Int[Tensor, "atoms"]]:
+    ) -> tuple[Int[Tensor, " atoms"], Int[Tensor, " atoms"]]:
         """Derive per-atom atom37 destination indices from an Atomworks atom array.
 
         For each atom in ``atom_array`` (the order the sampler's flat ``x_t``
@@ -636,7 +636,6 @@ def step(
         self,
         x_t: Float[Tensor, "batch atoms 3"],
         t: Float[Tensor, "*batch"] | float,
-        sigma_float: float,
         *,
         features: GenerativeModelInput[ProtpardelleConditioning] | None = None,
     ) -> Float[Tensor, "batch atoms 3"]:

tests/runs/test_runner.py

@@ -265,7 +265,8 @@ def test_dry_run_does_not_create_directories(
     """--dry-run prints commands but never touches the filesystem."""
     monkeypatch.setenv("HOME", str(tmp_path))
     results_dir = tmp_path / "results"
-    preset = loader.load_preset("rf3_partial", overrides=["jobs.0.gpu_count=1"])  # use 1 gpu so we don't need big nodes to test
+    # use 1 gpu so we don't need big nodes to test
+    preset = loader.load_preset("rf3_partial", overrides=["jobs.0.gpu_count=1"])
     runner.run(preset, results_dir=results_dir, dry_run=True)
     # results_dir gets created by run() (for log file location) but per-job
     # output subdirs must NOT exist after dry-run.