diff --git a/abogen/main.py b/abogen/main.py index 0b7875e..ad59430 100644 --- a/abogen/main.py +++ b/abogen/main.py @@ -10,6 +10,10 @@ from gui import abogen from utils import get_resource_path from constants import PROGRAM_NAME, VERSION +# Set environment variables for AMD ROCm +os.environ["MIOPEN_FIND_MODE"] = "FAST" +os.environ["MIOPEN_CONV_PRECISE_ROCM_TUNING"] = "0" + # Ensure sys.stdout and sys.stderr are valid in GUI mode if sys.stdout is None: sys.stdout = open(os.devnull, "w") diff --git a/abogen/utils.py b/abogen/utils.py index 07a8930..4be716c 100644 --- a/abogen/utils.py +++ b/abogen/utils.py @@ -105,7 +105,7 @@ def clean_text(text, *args, **kwargs): default_encoding = sys.getfilesystemencoding() -def create_process(cmd, stdin=None, text=True): +def create_process(cmd, stdin=None, text=True, capture_output=False): import logging logger = logging.getLogger(__name__) @@ -154,8 +154,8 @@ def create_process(cmd, stdin=None, text=True): proc = subprocess.Popen(cmd, **kwargs) - # Stream output to console in real-time - if proc.stdout: + # Stream output to console in real-time if not capturing + if proc.stdout and not capture_output: def _stream_output(stream): if text: # For text mode, read character by character for real-time output @@ -212,6 +212,33 @@ def calculate_text_length(text): return char_count +def get_gpu_vendor(): + os_name = platform.system() + cmd = None + if os_name == "Windows": + cmd = "wmic path win32_VideoController get name" + elif os_name == "Linux": + cmd = "lspci | grep -i 'vga\\|3d\\|display'" + + if cmd: + try: + # Call create_process with capture_output=True + proc = create_process(cmd, text=True, capture_output=True) + output, _ = proc.communicate(timeout=5) # Add a timeout + if proc.returncode == 0 and output: + output_lower = output.lower() + if "nvidia" in output_lower: + print("GPU Vendor: NVIDIA") + return "nvidia" + elif "amd" in output_lower or "radeon" in output_lower: + print("GPU Vendor: AMD") + return "amd" + except (subprocess.TimeoutExpired, FileNotFoundError, Exception): # Catch more specific exceptions + # Log error or handle as appropriate + return "unknown" # Fallback in case of any error + return "unknown" + + def get_gpu_acceleration(enabled): from torch.cuda import is_available