Fix/import vsa

[akaitsuki-ii]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @akaitsuki-ii, 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 focuses on architectural improvements and dependency management within the diffsynth_engine. It refactors the handling of attention-related parameters by moving the get_attn_kwargs logic directly into the base pipeline, ensuring that attention configurations are managed consistently across different pipeline implementations. Furthermore, it introduces conditional loading for the Video Sparse Attention (VSA) module, enhancing the system's modularity and reducing strict dependencies. The changes also include minor import adjustments and a more generalized approach to parallelization, contributing to a more robust and flexible codebase.

Highlights

• Refactored Attention Keyword Argument Handling: The get_attn_kwargs method, responsible for preparing parameters for attention mechanisms (like Sparge and Video Sparse Attention), has been moved from the AttentionConfig class to the BasePipeline class. This centralizes the logic for obtaining attention parameters within the pipeline itself, ensuring consistent management across different pipeline implementations.
• Conditional Video Sparse Attention (VSA) Import: The import of the vsa_core module, which provides Video Sparse Attention functionality, is now conditional on the VIDEO_SPARSE_ATTN_AVAILABLE flag. This change makes the VSA dependency optional, enhancing the system's modularity and flexibility.
• Generalized Parallelization Logic: The wrap_for_parallel function in diffsynth_engine/utils/parallel.py has been updated to use more generic attribute checks (hasattr and getattr) instead of specific class type checks (isinstance). This makes the parallelization wrapping mechanism more adaptable and less tightly coupled to specific model or pipeline types.
• Import and Export Adjustments: Various __init__.py files across the diffsynth_engine and diffsynth_engine/configs packages have been updated to correctly import and expose SpargeAttentionParams, VideoSparseAttentionParams, and LoraConfig in their new or adjusted locations, and LoraConfig was removed from the top-level __all__ export.

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 primarily focuses on refactoring. It makes the vsa library an optional dependency, which is a great improvement for usability. The logic for generating attention keyword arguments has been moved from the configuration classes to the base pipeline class, which is a better design. Additionally, the parallel execution utility has been decoupled from specific model classes, improving modularity.

My review includes a couple of suggestions for further improvement:

• Making a hardcoded patch_size for VSA configurable to improve flexibility.
• Restoring a type hint in the parallel utility to maintain code clarity.

[gemini-code-assist]

The patch_size tuple (1, 2, 2) is hardcoded in the call to get_vsa_kwargs. This could lead to issues if a different model using Video Sparse Attention requires a different patch size. It would be more robust to make this value configurable, for instance by passing it as an argument to get_attn_kwargs from the specific DiT model that is being used.