[model][NPU]: Z-image model support NPU

[Feng0w0]

NPU does not support the current operation result.append(self.freqs_cis[i][index]) at the RoPE of z_image_dit due to hardware reasons, and an additional NPU branch is added.

[gemini-code-assist]

Summary of Changes

Hello @Feng0w0, 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 NPU (Neural Processing Unit) compatibility for the z_image_dit model. It specifically addresses a hardware limitation in NPU environments where a particular tensor indexing operation within the RoPE calculation is not supported. By adding a conditional branch that uses torch.index_select when an NPU is detected, the model can now correctly perform its operations on NPU hardware, ensuring broader device support and functionality.

Highlights

• NPU Compatibility: Implemented conditional logic within the z_image_dit model's Rotary Positional Embedding (RoPE) calculation to ensure compatibility with NPU hardware.
• Hardware-Specific Operation: Replaced direct tensor indexing with torch.index_select for NPU environments, addressing a hardware limitation where direct indexing is not supported for a specific operation.
• New Import: Added an import for IS_NPU_AVAILABLE to dynamically detect NPU presence and enable the appropriate conditional logic.

🧠 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 adds NPU support for the Z-image DiT model by providing an alternative implementation for an indexing operation in the RopeEmbedder that is not supported on NPU hardware. The change is correct and necessary. I've suggested one improvement to make the device check more robust by inspecting the tensor's device directly, rather than relying on a global flag. This will prevent issues in environments with multiple available device types.

[gemini-code-assist]

Using the global IS_NPU_AVAILABLE flag can lead to incorrect behavior if the system has an NPU but the model is running on a different device (like CUDA or CPU). It's more robust to check the device of the tensor itself to decide which code path to take.

By checking ids.device.type, you ensure that the NPU-specific code path is only taken when the model is actually running on an NPU device. This change would also make the import of IS_NPU_AVAILABLE at the top of the file unnecessary.

Suggested change

	if IS_NPU_AVAILABLE:
	result.append(torch.index_select(self.freqs_cis[i], 0, index))
	else:
	result.append(self.freqs_cis[i][index])
	if ids.device.type == 'npu':
	result.append(torch.index_select(self.freqs_cis[i], 0, index))
	else:
	result.append(self.freqs_cis[i][index])