feat: reserve interface for other torch devices #27
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Reviewer's Guide by SourceryThis PR refactors device handling in the Final2x_core library by introducing a dedicated device utility function and updating the supported device list. The changes centralize device initialization logic and add support for DirectML devices. Class diagram for updated device handling in Final2x_coreclassDiagram
class CCRestoration {
-SRBaseModel _SR_class
+CCRestoration(SRConfig config)
+process(np.ndarray img) np.ndarray
}
class SRConfig {
+String device
+String pretrained_model_name
+String gh_proxy
}
class AutoModel {
+from_pretrained(String pretrained_model_name, Boolean fp16, Union<torch.device, String> device, String gh_proxy)
}
class device {
+get_device(String device) Union<torch.device, String>
}
CCRestoration --> SRConfig
CCRestoration --> AutoModel
AutoModel --> device
note for device "New utility function for device handling"
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
There was a problem hiding this comment.
Hey @Tohrusky - I've reviewed your changes and they look great!
Here's what I looked at during the review
- 🟡 General issues: 2 issues found
- 🟢 Security: all looks good
- 🟢 Testing: all looks good
- 🟢 Complexity: all looks good
- 🟢 Documentation: all looks good
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| elif device.startswith("xpu"): | ||
| return torch.device("xpu") | ||
| else: | ||
| print(f"Unknown device: {device}, use auto instead.") |
There was a problem hiding this comment.
suggestion: Consider using logger.warning instead of print for unknown device fallback
Using the logger would be more consistent with the rest of the codebase and provide better error tracking. Also consider raising a ValueError instead of falling back silently.
logger.warning(f"Unknown device: {device}, use auto instead.")| elif device.startswith("mps"): | ||
| return torch.device("mps") | ||
| elif device.startswith("directml"): | ||
| import torch_directml |
There was a problem hiding this comment.
suggestion: Add explicit error handling for optional directml import
Consider wrapping this in a try-except block to provide a more informative error message if torch_directml is not installed.
Suggested change
| import torch_directml | |
| try: | |
| import torch_directml | |
| except ImportError: | |
| raise ImportError("torch_directml is not installed. Please install it to use DirectML device.") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary by Sourcery
Add support for 'directml' device and refactor device initialization to improve flexibility in selecting torch devices. Update the project version to 3.0.2.
New Features:
Enhancements:
Chores: