📝 CodeRabbit Chat: Implement requested code changes

Author: coderabbitai[bot]

src/sampleworks/core/samplers/edm.py

@@ -137,8 +137,9 @@ def __post_init__(self) -> None:
 class AF3EDMSampler:
     """EDM-style sampler from AF3-like models.
 
-    All constants are configurable via constructor for model-specific values.
-    Default values match AF3 parameterization.
+    Initialized with a single :class:`EDMSamplerConfig` object that holds all
+    schedule hyperparameters and runtime options.  Default values in the config
+    match the AF3 parameterization.
 
     This sampler implements the EDM (Karras et al.) style sampling
     approach as used in AlphaFold3 and related models, which is the Euler
@@ -154,6 +155,18 @@ class AF3EDMSampler:
     """
 
     def __init__(self, config: EDMSamplerConfig) -> None:
+        """Initialize the sampler with a configuration object.
+
+        Parameters
+        ----------
+        config : EDMSamplerConfig
+            Configuration object containing all schedule hyperparameters
+            (``sigma_data``, ``s_max``, ``s_min``, ``p``, ``gamma_min``,
+            ``gamma_0``, ``noise_scale``, ``step_scale``) and runtime flags
+            (``augmentation``, ``align_to_input``,
+            ``alignment_reverse_diffusion``, ``scale_guidance_to_diffusion``,
+            ``device``).
+        """
         self.config = config
 
     def check_context(self, context: StepParams) -> None:
@@ -508,4 +521,4 @@ def step(
             denoised=x_hat_0_working_frame_t,
             loss=loss,
             log_proposal_correction=log_proposal_correction,
-        )
+        )

tests/conftest.py

@@ -253,7 +253,8 @@ def create_sampler_from_type(
         if device is not None:
             config_kwargs["device"] = device
         config = config_cls(**config_kwargs)
-        return create_component_from_info(info, config=config)
+        cls = _import_from_path(info.module_path)
+        return cls(config)
     return create_component_from_info(info, device=device, **extra_kwargs)
 
 
@@ -1067,4 +1068,4 @@ def perturbed_coords(
     torch.manual_seed(42)
     base = converging_mock_wrapper.target
     perturbation = torch.randn_like(base) * 0.1  # ty: ignore[invalid-argument-type]
-    return base, base + perturbation  # ty: ignore[invalid-return-type, unsupported-operator]
+    return base, base + perturbation  # ty: ignore[invalid-return-type, unsupported-operator]