[Diffusion] add taylorseer transformation

[legitnull]

1. Add example Wan-AI/Wan2.1-T2V-14B-Diffusers
2. Support saving multiple images/videos
3. Add module name/type matcher for Transformations
4. Add TimestepTrackerTransformation to update timesteps
5. Add TaylorSeerTransformation to apply TaylorSeer
   • Tested on Wan2.1 1.3B
6. Make hook.custom_forward explicit
7. States are now correctly reset between multiple runs

[gemini-code-assist]

Summary of Changes

Hello @legitnull, 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 a significant optimization for diffusion models by integrating the TaylorSeer transformation. This allows for faster inference by approximating module outputs during certain steps, reducing computational load. The changes involve a comprehensive update to the transformation framework, including a more flexible module selection mechanism and an enhanced runtime context to support the stateful nature of these optimizations. Configuration files have also been updated to leverage these new capabilities and support new model variants.

Highlights

• TaylorSeer Transformation: Introduced the TaylorSeer transformation, a novel technique to accelerate diffusion models by approximating module outputs using Taylor series. This includes new classes for managing Taylor series state and a custom hook to override module forward passes for approximation.
• Dynamic Transformation Application: Enhanced the inference engine to dynamically apply transformations to specific submodules based on configured targets, allowing for fine-grained control over where optimizations like TaylorSeer are applied.
• Runtime Context Enhancements: The RuntimeContext has been significantly updated to track diffusion timesteps and manage state scopes more robustly, which is crucial for stateful transformations like TaylorSeer.
• Flexible Module Selection: Implemented a new Selector mechanism within the transformation framework, enabling flexible targeting of modules for transformations using patterns (by name) or types (by class).
• Configuration Updates: Added a new inference configuration for a 1.3B model and updated an existing one (renamed from t2v.yaml to 14b.yaml) to incorporate the new transformation capabilities and model variants.
• Multiple Output Saving: The inference engine's save method now supports saving multiple image or video outputs, ensuring all generated content is properly stored.

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 introduces the TaylorSeer transformation, a significant performance optimization for diffusion models, by approximating some forward passes using Taylor series expansion. The implementation is well-structured, introducing a flexible transformation framework with module selectors, and adding necessary components like TimestepTrackerTransformation and StateScopeTransformation. The changes also include unit tests for the new transformations.

My review focuses on a few key areas for improvement:

• Correcting an inheritance issue in StateScopeHook.
• Removing leftover debug print statements.

Overall, this is a great contribution that enhances the inference capabilities. The code is clean and the new framework is well-designed.

[gemini-code-assist]

StateScopeHook inherits from ModelHook, but it doesn't call super().__init__() because the __init__ method has been removed. This means attributes from the base class (like fn_ref and _stateful) are not initialized. This is a potential bug and breaks the inheritance contract. Please add the __init__ method back to ensure proper initialization.

Suggested change

	def __init__(self) -> None:
	super().__init__()