[Diffusion] dit-precision refactor

[fsygd]

Motivation

Support convert_model_type in https://github.com/Wan-Video/Wan2.2/blob/main/generate.py#L220.
Found out https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/configs/pipeline_configs/base.py#L418 already did this. Just use sglang generate --dit-precision {fp32|bf16|fp16}

Modifications

Refactor for set_default_torch_dtype, make it more robust and remove redundant codes.

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @fsygd, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the mechanism for temporarily setting PyTorch's default data type. It consolidates the set_default_torch_dtype utility into a single, robust context manager and updates all call sites to use this centralized function. This change enhances the reliability of dtype management during model loading and reduces code duplication across the codebase.

Highlights

• Centralized Default Dtype Handling: The set_default_dtype context manager, previously duplicated or locally defined, has been centralized into python/sglang/multimodal_gen/runtime/loader/utils.py as set_default_torch_dtype.
• Improved Robustness: The set_default_torch_dtype context manager in utils.py has been made more robust by wrapping its yield statement in a try...finally block, ensuring the default dtype is always reset to its original value, even if errors occur within the with statement.
• Code Cleanup and Consistency: Redundant local implementations of the default dtype context manager have been removed, and all relevant modules now import and use the single, centralized utility function, improving code consistency and maintainability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the handling of the default torch dtype by moving the set_default_torch_dtype context manager to a central utility file and updating its usages. This improves code organization by removing redundant code. The change also correctly aims to make the context manager more robust by using a try...finally block. However, I've identified a critical bug in the new implementation where a recursive call would lead to a RecursionError. My review includes a suggestion to fix this issue.

[fsygd]

@mickqian @yhyang201 PTAL

[mickqian]

should we document it, in cli.md?

[mickqian]

/tag-and-rerun-ci

[fsygd]

should we document it, in cli.md?

ok, I will add it tonight.

[fsygd]

/rerun-failed-ci

[fsygd]

@mickqian could you please rerun ci? thx!

[mickqian]

/rerun-failed-ci

[fsygd]

/rerun-failed-ci