[NPU]:Replace 'cuda' in the project with abstract interfaces #1207

Feng0w0 · 2026-01-15T12:40:29Z

Replace the cuda field in the project with the abstract interface:
1.'cuda' -> get_device_type()
2.torch.cuda -> get_torch_device()

gemini-code-assist · 2026-01-15T12:40:58Z

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 significantly improves the project's hardware compatibility by abstracting device-specific code. It moves away from direct 'cuda' references to more general get_device_type() and get_torch_device() interfaces, making the codebase more flexible for different accelerators. A key part of this change is the implementation of a patch to ensure autocast functionality works correctly for NPU devices, thereby broadening the range of supported hardware without requiring extensive code modifications.

Highlights

Device Abstraction: Replaced hardcoded 'cuda' device references with abstract get_device_type() calls across numerous models and pipelines, enhancing device-agnosticism.
PyTorch CUDA Abstraction: Abstracted torch.cuda calls to get_torch_device() for retrieving the current device, further decoupling the codebase from CUDA-specific implementations.
NPU Autocast Patching: Introduced a specific patching mechanism for torch.amp.autocast to correctly handle torch.float32 on NPU devices, addressing a known compatibility issue.
Extended Compilation Support: Updated model compilation checks to explicitly include NPU devices, alongside CUDA, for improved hardware compatibility.

🧠 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.

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 effectively abstracts device-specific code by replacing hardcoded 'cuda' references to better support NPUs, primarily through helper functions like get_device_type() and get_torch_device(). The introduction of a monkey patch for torch.autocast is a clever solution for handling NPU-specific behavior. I've noted a couple of areas for improvement: there's some redundant code in diffsynth/models/longcat_video_dit.py where NPU-specific logic is manually added despite the monkey patch, and a potential bug in diffsynth/models/step1x_text_encoder.py with get_torch_device().current_device() that could fail on CPU-only systems. My suggestions aim to resolve these issues. Overall, this is a solid step towards making the codebase more device-agnostic.

gemini-code-assist · 2026-01-15T12:43:21Z

diffsynth/models/step1x_text_encoder.py

            self.model.config.hidden_size,
            dtype=torch.bfloat16,
-            device=torch.cuda.current_device(),
+            device=get_torch_device().current_device(),


Using get_torch_device().current_device() will cause a crash on CPU-only systems. The get_torch_device function falls back to torch.cuda when the device is 'cpu', and torch.cuda.current_device() will then fail if no CUDA device is available. A simpler and more robust approach is to use get_device_type() directly, as it returns a device string ('cpu', 'cuda', 'npu') that is accepted by PyTorch tensor creation functions.

Suggested change

device=get_torch_device().current_device(),

device=get_device_type(),

gemini-code-assist · 2026-01-15T12:43:22Z

diffsynth/models/step1x_text_encoder.py

            self.max_length,
            dtype=torch.long,
-            device=torch.cuda.current_device(),
+            device=get_torch_device().current_device(),


Using get_torch_device().current_device() will cause a crash on CPU-only systems. The get_torch_device function falls back to torch.cuda when the device is 'cpu', and torch.cuda.current_device() will then fail if no CUDA device is available. A simpler and more robust approach is to use get_device_type() directly, as it returns a device string ('cpu', 'cuda', 'npu') that is accepted by PyTorch tensor creation functions.

Suggested change

device=get_torch_device().current_device(),

device=get_device_type(),

gemini-code-assist · 2026-01-15T12:43:22Z

diffsynth/models/step1x_text_encoder.py

                (min(self.max_length, emb.shape[1] - 217)),
                dtype=torch.long,
-                device=torch.cuda.current_device(),
+                device=get_torch_device().current_device(),


Using get_torch_device().current_device() will cause a crash on CPU-only systems. The get_torch_device function falls back to torch.cuda when the device is 'cpu', and torch.cuda.current_device() will then fail if no CUDA device is available. A simpler and more robust approach is to use get_device_type() directly, as it returns a device string ('cpu', 'cuda', 'npu') that is accepted by PyTorch tensor creation functions.

Suggested change

device=get_torch_device().current_device(),

device=get_device_type(),

gemini-code-assist · 2026-01-15T12:43:22Z

diffsynth/models/longcat_video_dit.py

+            if IS_NPU_AVAILABLE:
+                torch.npu.set_autocast_enabled(True)


This conditional block is redundant. The monkey patch for torch.amp.autocast in diffsynth/core/npu_patch/npu_autocast_patch.py already handles this logic when get_device_type() returns 'npu'. Removing this will improve code clarity and avoid duplication. This comment applies to similar changes in this file (lines 590-591, 611-612, 626-627, and 810-811).

Feng0w0 and others added 8 commits January 9, 2026 18:06

[model][NPU]:Wan model rope use torch.complex64 in NPU

19ce304

[model][NPU]:Wan model rope use torch.complex64 in NPU

3b662da

Merge branch 'modelscope:main' into wan_rope

f4d06ce

[NPU]:Replace 'cuda' in the project with abstract interfaces

a3c2744

[NPU]:Replace 'cuda' in the project with abstract interfaces

209a350

Merge remote-tracking branch 'origin/cuda_replace' into cuda_replace

19e429d

[NPU]:Replace 'cuda' in the project with abstract interfaces

5c0b07d

[NPU]:Replace 'cuda' in the project with abstract interfaces

dce77ec

gemini-code-assist bot reviewed Jan 15, 2026

View reviewed changes

[NPU]:Replace 'cuda' in the project with abstract interfaces

ad91d41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[NPU]:Replace 'cuda' in the project with abstract interfaces #1207

[NPU]:Replace 'cuda' in the project with abstract interfaces #1207

Uh oh!

Feng0w0 commented Jan 15, 2026 •

edited

Loading

Uh oh!

gemini-code-assist bot commented Jan 15, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jan 15, 2026

Uh oh!

gemini-code-assist bot Jan 15, 2026

Uh oh!

gemini-code-assist bot Jan 15, 2026

Uh oh!

gemini-code-assist bot Jan 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	device=get_torch_device().current_device(),
	device=get_device_type(),

[NPU]:Replace 'cuda' in the project with abstract interfaces #1207

Are you sure you want to change the base?

[NPU]:Replace 'cuda' in the project with abstract interfaces #1207

Uh oh!

Conversation

Feng0w0 commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot commented Jan 15, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Feng0w0 commented Jan 15, 2026 •

edited

Loading