Feat/torch mps support #8
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds support for Apple Silicon MPS (Metal Performance Shaders) devices in addition to the existing CUDA support, enabling the AI-VIS library to run on Apple Silicon hardware. The changes centralize device selection and cache management in a new torch_utils module, replacing hard-coded CUDA-specific logic throughout the codebase.
Changes:
- Added a new
torch_utils.pymodule with centralized device selection (get_torch_device) and cache management (empty_torch_cache) functions - Refactored
RealESRGANerinitialization to use the new device selection utility instead of inline CUDA checks - Updated
RealESRGANModelto use device-agnostic tensor creation methods (.to(device)instead of.cuda()) - Replaced direct
torch.cuda.empty_cache()calls with the new backend-agnosticempty_torch_cache()function
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| aivis/torch_utils.py | New module providing centralized device selection and cache management with MPS support |
| aivis/models/realesrgan/utils.py | Refactored to use centralized device selection and removed redundant model assignment |
| aivis/models/realesrgan/models/realesrgan_model.py | Replaced .cuda() calls with .to(self.device) and added device parameter to tensor creation |
| aivis/aivis.py | Updated to use centralized device utilities instead of inline CUDA checks |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| return torch.device("cpu") | ||
|
|
||
| if torch.cuda.is_available(): | ||
| return torch.device("cuda") |
There was a problem hiding this comment.
The function doesn't handle specific CUDA device IDs. When gpu_id is set to a specific device like "0" or "1", the function returns torch.device("cuda") without the device index, which may not match the intended device. The function should return torch.device(f"cuda:{gpu_id}") when CUDA is available and a specific gpu_id is provided.
| return torch.device("cuda") | |
| cuda_device_str = "cuda" if gpu_id is None else f"cuda:{gpu_id}" | |
| return torch.device(cuda_device_str) |
| model.eval() | ||
| self.model = model.to(self.device) | ||
| self.model.eval() | ||
| self.model = self.model.to(self.device) |
There was a problem hiding this comment.
The model is being reassigned after calling .to(self.device). The assignment is unnecessary since .to() is an in-place operation for module parameters. Line 71 should just be self.model.to(self.device) without the assignment, or keep the assignment but remove the redundancy with line 53.
| self.model = self.model.to(self.device) | |
| self.model.to(self.device) |
| try: | ||
| mps.empty_cache() | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
The broad exception handler silently swallows all exceptions without logging. Consider either narrowing the exception type or adding logging to help diagnose issues with MPS cache clearing, especially since this is a new feature.
|
|
||
| if model is not None: | ||
| self.model = model.to(self.device) | ||
| self.model = model |
There was a problem hiding this comment.
This assignment at line 53 is redundant since the model is reassigned at line 71 after loading state dict and moving to device. The assignment at line 53 should be removed as the model needs to be moved to the target device after state dict loading anyway.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78d5404320
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if gpu_id is None or str(gpu_id) == "-1": | ||
| return torch.device("cpu") |
There was a problem hiding this comment.
The new get_torch_device treats gpu_id=None as a hard disable and immediately returns CPU, which means callers that omit gpu_id (e.g., RealESRGANer defaults it to None) will never use CUDA even when it is available. Previously, the default path selected CUDA automatically, so this is a regression that can silently force CPU inference and significantly slow down workloads unless every caller now passes a non-None gpu_id.
Useful? React with 👍 / 👎.
1 participant
No description provided.