support Wan/fp8

[sir1st-inc]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @sir1st-inc, 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 introduces support for FP8 (float8) precision within the WAN models, specifically targeting the DiT and T5 text encoder components. The changes involve refining data type handling in the time embedding calculation to ensure consistency and, more significantly, implementing conditional logic to enable FP8 autocasting within the pipeline's from_pretrained method. This enhancement aims to leverage lower precision floating-point numbers for potential improvements in performance and memory efficiency.

Highlights

• Type Consistency in Time Embedding: The sinusoidal_embedding_1d output is now explicitly cast to x.dtype before being passed to self.time_embedding. This ensures type consistency within the time embedding calculation.
• Flexible Timestep Type Handling: The explicit dtype conversion for the timestep tensor has been removed, allowing it to retain its original data type when moved to the device. This change likely facilitates the integration of mixed-precision operations.
• FP8 Autocasting for DiT Model: New logic has been introduced to enable FP8 (float8) autocasting for the 'dit' model. If config.model_dtype is torch.float8_e4m3fn, the pipeline's compute dtype is set to torch.bfloat16, and enable_fp8_autocast is called for the 'dit' model.
• FP8 Autocasting for Text Encoder: Similar to the 'dit' model, FP8 autocasting is now supported for the 'text_encoder'. If config.t5_dtype is torch.float8_e4m3fn, the compute dtype is set to torch.bfloat16, and enable_fp8_autocast is applied to the 'text_encoder'.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 adds support for fp8 in the Wan pipeline. The changes involve updating data type conversions for time embeddings and enabling fp8 autocasting for the DiT and text encoder models. My review includes a suggestion to refactor duplicated code for enabling fp8 autocast to improve maintainability, and a minor style fix to adhere to PEP 8.

[gemini-code-assist]

For better readability and adherence to the PEP 8 style guide, it's recommended to add a space around the assignment operator.¹

Suggested change

	t=sinusoidal_embedding_1d(self.freq_dim, timestep).to(x.dtype)
	t = sinusoidal_embedding_1d(self.freq_dim, timestep).to(x.dtype)

Style Guide References

Footnotes

1. PEP 8 E225 recommends a single space around most operators, including assignment (=), to improve code readability. ↩

[gemini-code-assist]

The logic for enabling fp8 autocast is duplicated for dit and text_encoder. This can be refactored to reduce code duplication and improve maintainability. You can collect the model names that require fp8 and then call enable_fp8_autocast once.

Suggested change

	if config.model_dtype == torch.float8_e4m3fn:
	pipe.dtype = torch.bfloat16 # compute dtype
	pipe.enable_fp8_autocast(
	model_names=["dit"], compute_dtype=pipe.dtype, use_fp8_linear=config.use_fp8_linear
	)
	if config.t5_dtype == torch.float8_e4m3fn:
	pipe.dtype = torch.bfloat16 # compute dtype
	pipe.enable_fp8_autocast(
	model_names=["text_encoder"], compute_dtype=pipe.dtype, use_fp8_linear=config.use_fp8_linear
	)
	fp8_models = []
	if config.model_dtype == torch.float8_e4m3fn:
	fp8_models.append("dit")
	if config.t5_dtype == torch.float8_e4m3fn:
	fp8_models.append("text_encoder")
	if fp8_models:
	pipe.dtype = torch.bfloat16 # compute dtype
	pipe.enable_fp8_autocast(
	model_names=fp8_models, compute_dtype=pipe.dtype, use_fp8_linear=config.use_fp8_linear
	)