Skip to content

Conversation

@delmalih
Copy link

@delmalih delmalih commented Nov 10, 2025

What does this PR do?

This PR improves the documentation and type hints in scheduling_amused.py to follow the project conventions as described in #9567.

Changes made:

  • Helper functions:

    • Added complete type hints and docstrings to gumbel_noise() function
    • Added complete type hints and docstrings to mask_by_random_topk() function
  • AmusedSchedulerOutput class:

    • Fixed type hint for pred_original_sample attribute (torch.TensorOptional[torch.Tensor])
  • AmusedScheduler class:

    • Added type hints for class attributes (temperatures and timesteps)
    • Added comprehensive docstring to __init__() method
    • Added comprehensive docstring to set_timesteps() method with proper return type annotation
    • Added comprehensive docstring to step() method with detailed parameter descriptions
    • Fixed parameter type hints in step(): timestep and starting_mask_ratio
    • Added complete type hints and docstring to add_noise() method

All docstrings follow the standard format with proper Args and Returns sections, using the backtick notation for types as per project conventions.

Context

This is part of a broader effort to improve documentation across all scheduler files in src/diffusers/schedulers/. See issue #9567 for details.

Related PRs:

Testing

  • make style passed (with automatic formatting applied)
  • make quality passed
  • Module imports successfully without errors
  • Basic functionality tests passed for all modified functions

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline?
  • Did you read our philosophy doc (important for complex PRs)?
  • Was this discussed/approved via a GitHub issue or the forum? issue.
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests? (N/A - documentation only changes)

Who can review?

@yiyixuxu (Schedulers maintainer)

@delmalih delmalih marked this pull request as draft November 10, 2025 16:00
@delmalih delmalih marked this pull request as ready for review November 10, 2025 16:01
@delmalih delmalih force-pushed the improve-docstring-scheduling-amused branch 2 times, most recently from 8c8c848 to 67f931b Compare November 10, 2025 16:07
- Add complete type hints for helper functions (gumbel_noise, mask_by_random_topk)
- Enhance AmusedSchedulerOutput with proper Optional typing
- Add comprehensive docstrings for AmusedScheduler class
- Improve __init__, set_timesteps, step, and add_noise methods
- Fix type hints to match documentation conventions
- All changes follow project standards from issue huggingface#9567
- Update type hints for `prev_sample` and `pred_original_sample` in `AmusedSchedulerOutput` to reflect their tensor types.
- Improve docstring for `gumbel_noise` to specify the output tensor's dtype and device.
- Refine `AmusedScheduler` class documentation, including detailed descriptions of the masking schedule and temperature parameters.
- Adjust type hints in `set_timesteps` and `step` methods for better clarity and consistency.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant