Fix diffusion inferers to support diffusers-style schedulers#8870
Open
ugbotueferhire wants to merge 4 commits into
Open
Fix diffusion inferers to support diffusers-style schedulers#8870ugbotueferhire wants to merge 4 commits into
ugbotueferhire wants to merge 4 commits into
Conversation
ugbotueferhire
requested review from
KumoLiu,
Nic-Ma,
ericspod and
virginiafdez
as code owners
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR standardizes scheduler integration by adding inspect-based helpers that normalize scheduler.step arguments/outputs and expose scheduler config values. DiffusionInferer.sample and ControlNetDiffusionInferer.sample use a unified _scheduler_step for next-sample computation. get_likelihood now uses the passed scheduler's add_noise, reads variance_type/prediction_type/clip settings from scheduler config, and computes posterior mean/variance locally. Tests add diffusers-gated cases using DDPMScheduler to validate sampling and likelihood intermediates. Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (6)
monai/inferers/inferer.py (1)
865-900: ⚡ Quick winAdd docstrings for new helper methods.
These methods lack docstrings. Per coding guidelines, all definitions should have Google-style docstrings describing parameters, return values, and raised exceptions.
📝 Suggested docstrings
`@staticmethod` def _scheduler_step_supports_kwarg(scheduler: Scheduler, kwarg: str) -> bool: + """Check if scheduler.step accepts a specific keyword argument. + + Args: + scheduler: The scheduler instance to inspect. + kwarg: Name of the keyword argument to check. + + Returns: + True if the kwarg is in scheduler.step's signature, False otherwise. + """ try: return kwarg in inspect.signature(scheduler.step).parameters except (TypeError, ValueError): return False `@staticmethod` def _get_previous_sample_from_step_output(step_output: Any) -> torch.Tensor: + """Extract the previous sample tensor from scheduler.step output. + + Args: + step_output: Output from scheduler.step, which may be a tuple, + Mapping, or object with prev_sample attribute. + + Returns: + The previous sample tensor. + + Raises: + TypeError: If step_output format is unsupported. + """ if isinstance(step_output, tuple): return step_output[0] if isinstance(step_output, Mapping): return step_output["prev_sample"] if hasattr(step_output, "prev_sample"): return step_output.prev_sample - raise TypeError("Unsupported scheduler.step output. Expected a tuple or an object with `prev_sample`.") + raise TypeError("Unsupported scheduler.step output. Expected a tuple, Mapping, or object with `prev_sample`.") def _scheduler_step( self, scheduler: Scheduler, model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, next_timestep: int | torch.Tensor | None = None, ) -> torch.Tensor: + """Perform a scheduler step with unified handling for various scheduler types. + + Args: + scheduler: The diffusion scheduler. + model_output: Output from the diffusion model. + timestep: Current timestep. + sample: Current sample tensor. + next_timestep: Next timestep (used only for RFlowScheduler). + + Returns: + The previous sample tensor from the scheduler step. + """ step_kwargs = {}As per coding guidelines: "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/inferers/inferer.py` around lines 865 - 900, Add Google-style docstrings to the three helper methods _scheduler_step_supports_kwarg, _get_previous_sample_from_step_output, and _scheduler_step: for _scheduler_step_supports_kwarg describe the parameters (scheduler: Scheduler, kwarg: str), the boolean return value and that it may return False if inspect.signature raises TypeError/ValueError; for _get_previous_sample_from_step_output document the parameter (step_output: Any), the torch.Tensor return (the previous sample) and the TypeError raised for unsupported outputs; for _scheduler_step document parameters (scheduler, model_output, timestep, sample, next_timestep), the torch.Tensor return value (the previous sample), and mention that it relies on _scheduler_step_supports_kwarg and _get_previous_sample_from_step_output and may raise the TypeError propagated by the latter; place each docstring immediately above its respective def using Google-style sections Args, Returns, and Raises.tests/inferers/test_latent_diffusion_inferer.py (3)
444-474: ⚡ Quick winAdd a docstring to the new latent diffusers compatibility test.
Line 446 defines a new test method without a docstring. Add a short Google-style docstring describing the scenario and expected output contract.
As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/inferers/test_latent_diffusion_inferer.py` around lines 444 - 474, Add a Google-style docstring to the test method test_diffusers_style_ddpm_sample_shape describing the scenario (compatibility sampling with Diffusers-style DDPM scheduler on latent space), the parameters used (ae_model_type, dm_model_type, input_noise, autoencoder_model, diffusion_model, scheduler), the expected return (tensor of shape equal to input_shape) and any exceptions (e.g., device availability). Place the docstring immediately below the def line, using short summary, Args (list the key args used in the test), and Returns (describe the expected tensor shape) sections so the test clearly documents its behavior and contract.
316-316: ⚡ Quick winBroaden
TEST_CASES_DIFFUSERScoverage.Line 316 limits the new compatibility test to one fixture. Add at least one more representative case (e.g., a non-
AutoencoderKLpath or 3D path) to reduce regression risk on scheduler-output normalization.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/inferers/test_latent_diffusion_inferer.py` at line 316, TEST_CASES_DIFFUSERS currently contains only a single entry (TEST_CASES[0]), which limits coverage; update its definition to include at least one additional representative test case from the TEST_CASES array (for example an entry that exercises a non-AutoencoderKL path or a 3D path) so the compatibility test covers both diffusion/scheduler normalization and a different decoding path. Locate TEST_CASES_DIFFUSERS and replace the single-element list with a list of two or more TEST_CASES entries (e.g., TEST_CASES[0] and TEST_CASES[1] or another index whose fixture represents the non-AutoencoderKL/3D scenario) so the tests exercise both code paths.
319-342: ⚡ Quick winAdd Google-style docstrings to the new scheduler wrapper definitions.
Line 319 and Line 326 add new definitions without docstrings. Please document args, returns, and raised exceptions.
As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/inferers/test_latent_diffusion_inferer.py` around lines 319 - 342, Add Google-style docstrings to the new classes and method: document DiffusersLikeSchedulerOutput (describe prev_sample and pred_original_sample types and purpose) and DiffusersStyleDDPMScheduler.step (document args: model_output, timestep, sample, generator, return_dict; the return value when return_dict is True vs False; and any exceptions that may be raised, e.g., from the super().step call). Place the class-level docstring above DiffusersLikeSchedulerOutput and the method docstring on DiffusersStyleDDPMScheduler.step following Google style sections: Args, Returns, Raises, and keep descriptions brief and type-annotated to match the signatures.tests/inferers/test_diffusion_inferer.py (2)
154-170: ⚡ Quick winAdd a docstring to the new test method.
Line 156 introduces a new test definition without a docstring. Add a short Google-style docstring stating the scenario and expected behavior.
As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/inferers/test_diffusion_inferer.py` around lines 154 - 170, Add a Google-style docstring to the test_diffusers_style_ddpm_sampler method describing the scenario (sampling from DiffusionModelUNet with DiffusersStyleDDPMScheduler using 10 inference steps and saving intermediates) and the expected behavior (output sample has same shape as input noise and intermediates length equals number of timesteps); place the docstring immediately below the def test_diffusers_style_ddpm_sampler(...) line and include brief sections for Args (model_params, input_shape), Returns (None), and Expected behavior/assertions (sample shape and intermediates length).
58-80: ⚡ Quick winAdd Google-style docstrings to the new scheduler shim definitions.
Line 58 and Line 65 add new definitions without docstrings. Please add docstrings documenting args, returns, and raised exceptions.
As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/inferers/test_diffusion_inferer.py` around lines 58 - 80, Add Google-style docstrings for the new classes and methods: document DiffusersLikeSchedulerOutput.__init__ parameters prev_sample and pred_original_sample and that it stores them as attributes, and document DiffusersStyleDDPMScheduler.step parameters model_output, timestep, sample, generator, and return_dict, describe the return types (either DiffusersLikeSchedulerOutput when return_dict is True or Tuple[torch.Tensor, torch.Tensor] otherwise) and note any exceptions propagated from the parent DDPMScheduler.step; place the docstrings immediately above the class DiffusersLikeSchedulerOutput and the method DiffusersStyleDDPMScheduler.step using Google-style sections Args, Returns, and Raises.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@monai/inferers/inferer.py`:
- Around line 865-900: Add Google-style docstrings to the three helper methods
_scheduler_step_supports_kwarg, _get_previous_sample_from_step_output, and
_scheduler_step: for _scheduler_step_supports_kwarg describe the parameters
(scheduler: Scheduler, kwarg: str), the boolean return value and that it may
return False if inspect.signature raises TypeError/ValueError; for
_get_previous_sample_from_step_output document the parameter (step_output: Any),
the torch.Tensor return (the previous sample) and the TypeError raised for
unsupported outputs; for _scheduler_step document parameters (scheduler,
model_output, timestep, sample, next_timestep), the torch.Tensor return value
(the previous sample), and mention that it relies on
_scheduler_step_supports_kwarg and _get_previous_sample_from_step_output and may
raise the TypeError propagated by the latter; place each docstring immediately
above its respective def using Google-style sections Args, Returns, and Raises.
In `@tests/inferers/test_diffusion_inferer.py`:
- Around line 154-170: Add a Google-style docstring to the
test_diffusers_style_ddpm_sampler method describing the scenario (sampling from
DiffusionModelUNet with DiffusersStyleDDPMScheduler using 10 inference steps and
saving intermediates) and the expected behavior (output sample has same shape as
input noise and intermediates length equals number of timesteps); place the
docstring immediately below the def test_diffusers_style_ddpm_sampler(...) line
and include brief sections for Args (model_params, input_shape), Returns (None),
and Expected behavior/assertions (sample shape and intermediates length).
- Around line 58-80: Add Google-style docstrings for the new classes and
methods: document DiffusersLikeSchedulerOutput.__init__ parameters prev_sample
and pred_original_sample and that it stores them as attributes, and document
DiffusersStyleDDPMScheduler.step parameters model_output, timestep, sample,
generator, and return_dict, describe the return types (either
DiffusersLikeSchedulerOutput when return_dict is True or Tuple[torch.Tensor,
torch.Tensor] otherwise) and note any exceptions propagated from the parent
DDPMScheduler.step; place the docstrings immediately above the class
DiffusersLikeSchedulerOutput and the method DiffusersStyleDDPMScheduler.step
using Google-style sections Args, Returns, and Raises.
In `@tests/inferers/test_latent_diffusion_inferer.py`:
- Around line 444-474: Add a Google-style docstring to the test method
test_diffusers_style_ddpm_sample_shape describing the scenario (compatibility
sampling with Diffusers-style DDPM scheduler on latent space), the parameters
used (ae_model_type, dm_model_type, input_noise, autoencoder_model,
diffusion_model, scheduler), the expected return (tensor of shape equal to
input_shape) and any exceptions (e.g., device availability). Place the docstring
immediately below the def line, using short summary, Args (list the key args
used in the test), and Returns (describe the expected tensor shape) sections so
the test clearly documents its behavior and contract.
- Line 316: TEST_CASES_DIFFUSERS currently contains only a single entry
(TEST_CASES[0]), which limits coverage; update its definition to include at
least one additional representative test case from the TEST_CASES array (for
example an entry that exercises a non-AutoencoderKL path or a 3D path) so the
compatibility test covers both diffusion/scheduler normalization and a different
decoding path. Locate TEST_CASES_DIFFUSERS and replace the single-element list
with a list of two or more TEST_CASES entries (e.g., TEST_CASES[0] and
TEST_CASES[1] or another index whose fixture represents the non-AutoencoderKL/3D
scenario) so the tests exercise both code paths.
- Around line 319-342: Add Google-style docstrings to the new classes and
method: document DiffusersLikeSchedulerOutput (describe prev_sample and
pred_original_sample types and purpose) and DiffusersStyleDDPMScheduler.step
(document args: model_output, timestep, sample, generator, return_dict; the
return value when return_dict is True vs False; and any exceptions that may be
raised, e.g., from the super().step call). Place the class-level docstring above
DiffusersLikeSchedulerOutput and the method docstring on
DiffusersStyleDDPMScheduler.step following Google style sections: Args, Returns,
Raises, and keep descriptions brief and type-annotated to match the signatures.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 2f006d49-4b91-4a20-9ae2-763bdb819a6f
📒 Files selected for processing (3)
monai/inferers/inferer.pytests/inferers/test_diffusion_inferer.pytests/inferers/test_latent_diffusion_inferer.py
Member
|
@virginiafdez could you please have a look at this one? Thanks! |
virginiafdez
requested changes
May 28, 2026
Contributor
There was a problem hiding this comment.
- The changes to
infererlook good and make sense. - I would make a modification on the tests to make the changes less disruptive.
- In addition, since the main point of this issue was to make MONAI inferer class compatible with the diffusers library, I would add a method on both testing files that queries an example scheduler from the
diffuserslibrary.
e.g.
from diffusers import DDPMScheduler as DiffusersDDPMScheduler
model = DiffusionModelUNet(
spatial_dims=2,
in_channels=1,
out_channels=1,
channels=[32, 64],
attention_levels=[False, True],
num_res_blocks=1,
num_head_channels=32,
).to(device)
scheduler = DiffusersDDPMScheduler(
num_train_timesteps=1000,
beta_schedule="linear",
prediction_type="epsilon",
)
scheduler.set_timesteps(num_inference_steps=50)
inferer = DiffusionInferer(scheduler=scheduler)
batch_size = 2
image_size = 32
inputs = torch.randn(batch_size, 1, image_size, image_size).to(device)
noise = torch.randn_like(inputs)
timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (batch_size,)).long().to(device)
with torch.no_grad():
prediction = inferer(
inputs=inputs,
diffusion_model=model,
noise=noise,
timesteps=timesteps
)
Another test: what happens if you import DDPMScheduler from the diffusers library and call get_likelihood? Will it work? Or something needs to do with renaming
Comment on lines
+58
to
+82
| class DiffusersLikeSchedulerOutput: | ||
| def __init__(self, prev_sample: torch.Tensor, pred_original_sample: torch.Tensor) -> None: | ||
| self.prev_sample = prev_sample | ||
| self.pred_original_sample = pred_original_sample | ||
|
|
||
|
|
||
| class DiffusersStyleDDPMScheduler(DDPMScheduler): | ||
| def step( | ||
| self, | ||
| model_output: torch.Tensor, | ||
| timestep: int, | ||
| sample: torch.Tensor, | ||
| generator: torch.Generator | None = None, | ||
| return_dict: bool = True, | ||
| ): | ||
| prev_sample, pred_original_sample = super().step( | ||
| model_output=model_output, timestep=timestep, sample=sample, generator=generator | ||
| ) | ||
| if return_dict: | ||
| return DiffusersLikeSchedulerOutput( | ||
| prev_sample=prev_sample, pred_original_sample=pred_original_sample | ||
| ) | ||
| return prev_sample, pred_original_sample | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
For a single method, I don't think this class is necessary. We can have this as a method on the main test class. But not 100% sure @ericspod
Member
There was a problem hiding this comment.
It hasn't been our style or practice with MONAI to create output classes like this, which here are essentially like namedtuple or dataclass. Was there a specific reason for this usage?
Comment on lines
+319
to
+342
| class DiffusersLikeSchedulerOutput: | ||
| def __init__(self, prev_sample: torch.Tensor, pred_original_sample: torch.Tensor) -> None: | ||
| self.prev_sample = prev_sample | ||
| self.pred_original_sample = pred_original_sample | ||
|
|
||
|
|
||
| class DiffusersStyleDDPMScheduler(DDPMScheduler): | ||
| def step( | ||
| self, | ||
| model_output: torch.Tensor, | ||
| timestep: int, | ||
| sample: torch.Tensor, | ||
| generator: torch.Generator | None = None, | ||
| return_dict: bool = True, | ||
| ): | ||
| prev_sample, pred_original_sample = super().step( | ||
| model_output=model_output, timestep=timestep, sample=sample, generator=generator | ||
| ) | ||
| if return_dict: | ||
| return DiffusersLikeSchedulerOutput( | ||
| prev_sample=prev_sample, pred_original_sample=pred_original_sample | ||
| ) | ||
| return prev_sample, pred_original_sample | ||
|
|
Contributor
There was a problem hiding this comment.
Same as comment for test_diffusion_inferer.
Author
|
Hi @virginiafdez @ericspod , thanks for the review. I removed the custom diffusers-style scheduler test shims and replaced them with optional tests using the real I also addressed the CodeRabbit comment by adding Verified with: python -m ruff check monai/inferers/inferer.py tests/inferers/test_diffusion_inferer.py tests/inferers/test_latent_diffusion_inferer.py
python -m pytest tests/inferers/test_diffusion_inferer.py -q
python -m pytest tests/inferers/test_latent_diffusion_inferer.py -q |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
monai/inferers/inferer.py (1)
1157-1199:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFix
learned_rangevariance handling + KL variance logging
predicted_varianceis split frommodel_outputforvariance_type in ["learned","learned_range"], but forlearned_rangeit’s an interpolation signal (expected in[-1, 1]), not a variance tensor; KL currently doestorch.log(predicted_variance)atmonai/inferers/inferer.py(~1198-1199, also ~1828-1829), which can produce NaNs and wrong KL terms._get_posterior_variance(..., variance_type="learned_range")interpolates usingmin_log = variance/max_log = scheduler.betas[t]withoutlog/exp(so it interpolates in variance-space despite the “log” naming), which disagrees with the learned_range formulation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/inferers/inferer.py` around lines 1157 - 1199, predicted_variance is treated as a raw variance for all variance_type values which breaks learned_range (it is an interpolation signal in [-1,1]), and torch.log(predicted_variance) produces NaNs; update _get_posterior_variance and callers to handle "learned_range" by first converting the learned_range signal into log-variance space (e.g. map signal in [-1,1] to a interpolation weight and interpolate between log(min_var) and log(max_var) or take logs of endpoints then interpolate) instead of interpolating raw variances, and change the computation of log_predicted_variance to compute the log from that converted variance (not directly torch.log of the learned_range tensor); specifically modify _get_posterior_variance(scheduler, timestep, predicted_variance, variance_type) and the usage sites of predicted_variance/log_predicted_variance in the infer loop (symbols: predicted_variance, variance_type, _get_posterior_variance, log_predicted_variance) so learned_range yields a valid positive variance and a stable log-variance for KL.
🧹 Nitpick comments (1)
monai/inferers/inferer.py (1)
865-935: ⚡ Quick winAdd docstrings to the new scheduler helpers.
These helpers are the new compatibility contract, but they ship without the required Google-style docstrings. Please document args, returns, and raised errors for each added definition.
As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/inferers/inferer.py` around lines 865 - 935, Add Google-style docstrings to the new scheduler helper methods: _scheduler_step_supports_kwarg, _get_previous_sample_from_step_output, _get_scheduler_name, _get_scheduler_config_value, _get_posterior_mean, and _get_posterior_variance; for each docstring include a short description, Args (types and purpose: e.g., scheduler: Scheduler, kwarg: str, timestep: int|torch.Tensor, x_0/x_t/predicted_variance: torch.Tensor, name/default for config), Returns (type and meaning), and Raises (any TypeError/ValueError conditions like unsupported step output or invalid inputs), and mention any special behavior (e.g., handling of Mapping vs attributes, learned variance branch) so callers know expected outputs and error conditions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/inferers/test_diffusion_inferer.py`:
- Around line 130-160: The test currently exercises DiffusionInferer.__call__
but never executes the scheduler-step sampling path; update
test_diffusers_ddpm_call to call DiffusionInferer.sample (or add a new test that
calls sample) with the same setup (scheduler=DiffusersDDPMScheduler,
diffusion_model=DiffusionModelUNet, inputs, noise, timesteps) so the
scheduler.step path is exercised; ensure the returned tensor shape is asserted
(self.assertEqual(prediction.shape, inputs.shape)) and keep device,
model.to(device), scheduler.set_timesteps(...) and torch.no_grad() usage intact
so the sampling branch and scheduler-step compatibility are covered.
---
Outside diff comments:
In `@monai/inferers/inferer.py`:
- Around line 1157-1199: predicted_variance is treated as a raw variance for all
variance_type values which breaks learned_range (it is an interpolation signal
in [-1,1]), and torch.log(predicted_variance) produces NaNs; update
_get_posterior_variance and callers to handle "learned_range" by first
converting the learned_range signal into log-variance space (e.g. map signal in
[-1,1] to a interpolation weight and interpolate between log(min_var) and
log(max_var) or take logs of endpoints then interpolate) instead of
interpolating raw variances, and change the computation of
log_predicted_variance to compute the log from that converted variance (not
directly torch.log of the learned_range tensor); specifically modify
_get_posterior_variance(scheduler, timestep, predicted_variance, variance_type)
and the usage sites of predicted_variance/log_predicted_variance in the infer
loop (symbols: predicted_variance, variance_type, _get_posterior_variance,
log_predicted_variance) so learned_range yields a valid positive variance and a
stable log-variance for KL.
---
Nitpick comments:
In `@monai/inferers/inferer.py`:
- Around line 865-935: Add Google-style docstrings to the new scheduler helper
methods: _scheduler_step_supports_kwarg, _get_previous_sample_from_step_output,
_get_scheduler_name, _get_scheduler_config_value, _get_posterior_mean, and
_get_posterior_variance; for each docstring include a short description, Args
(types and purpose: e.g., scheduler: Scheduler, kwarg: str, timestep:
int|torch.Tensor, x_0/x_t/predicted_variance: torch.Tensor, name/default for
config), Returns (type and meaning), and Raises (any TypeError/ValueError
conditions like unsupported step output or invalid inputs), and mention any
special behavior (e.g., handling of Mapping vs attributes, learned variance
branch) so callers know expected outputs and error conditions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b712c4b5-e829-4a54-b6b2-11575813ce23
📒 Files selected for processing (3)
monai/inferers/inferer.pytests/inferers/test_diffusion_inferer.pytests/inferers/test_latent_diffusion_inferer.py
Signed-off-by: ugbotueferhire <ugbotueferhire@gmail.com>
I, ugbotueferhire <ugbotueferhire@gmail.com>, hereby add my Signed-off-by to this commit: 8f95fb1 Signed-off-by: ugbotueferhire <ugbotueferhire@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #8597.
Description
This PR adds compatibility for
diffusers-style schedulers in MONAI's diffusion inferers (DiffusionInfererandLatentDiffusionInferer).Specifically, this includes:
inferer.pynow detectsreturn_dictsupport onscheduler.step, requests tuple output when available, and normalizes both MONAI tuple returns anddiffusers-styleprev_sampleoutputs. This shared helper also benefits the ControlNet diffusion inferers.schedulerinstance passed into the call.test_diffusion_inferer.pyandtest_latent_diffusion_inferer.pyusing a small diffusers-style DDPM shim that returns aprev_sampleoutput object by default. Verified locally (31 passedand46 passed).Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.