227 lines
7.6 KiB
Python
227 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Hardware Detection and Optimization Module
|
|
Detects available hardware (GPU, CPU) and recommends optimal encoding settings.
|
|
"""
|
|
|
|
import subprocess
|
|
import platform
|
|
import json
|
|
from pathlib import Path
|
|
from .logger_helper import setup_logger
|
|
|
|
logger = setup_logger(Path(__file__).parent.parent / "logs")
|
|
|
|
|
|
class HardwareInfo:
|
|
"""Detects and stores hardware capabilities."""
|
|
|
|
def __init__(self):
|
|
self.platform_name = platform.system() # Windows, Linux, Darwin
|
|
self.cpu_cores = self._get_cpu_cores()
|
|
self.gpu_type = self._detect_gpu()
|
|
self.available_encoders = self._check_ffmpeg_encoders()
|
|
self.recommended_encoder = self._recommend_encoder()
|
|
self.recommended_settings = self._get_recommended_settings()
|
|
|
|
def _get_cpu_cores(self):
|
|
"""Get number of CPU cores."""
|
|
import multiprocessing
|
|
return multiprocessing.cpu_count()
|
|
|
|
def _detect_gpu(self):
|
|
"""Detect GPU type (NVIDIA, AMD, Intel, None)."""
|
|
if self.platform_name == "Windows":
|
|
return self._detect_gpu_windows()
|
|
elif self.platform_name == "Linux":
|
|
return self._detect_gpu_linux()
|
|
else:
|
|
return None
|
|
|
|
def _detect_gpu_windows(self):
|
|
"""Detect GPU on Windows using DXDIAG or WMI."""
|
|
try:
|
|
# Try using wmic (Windows only)
|
|
result = subprocess.run(
|
|
["wmic", "path", "win32_videocontroller", "get", "name"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5
|
|
)
|
|
gpu_info = result.stdout.lower()
|
|
|
|
if "nvidia" in gpu_info or "geforce" in gpu_info or "quadro" in gpu_info:
|
|
return "nvidia"
|
|
elif "amd" in gpu_info or "radeon" in gpu_info:
|
|
return "amd"
|
|
elif "intel" in gpu_info:
|
|
return "intel"
|
|
except Exception as e:
|
|
logger.warning(f"Could not detect GPU via WMI: {e}")
|
|
|
|
return None
|
|
|
|
def _detect_gpu_linux(self):
|
|
"""Detect GPU on Linux using lspci."""
|
|
try:
|
|
result = subprocess.run(
|
|
["lspci"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5
|
|
)
|
|
gpu_info = result.stdout.lower()
|
|
|
|
if "nvidia" in gpu_info:
|
|
return "nvidia"
|
|
elif "amd" in gpu_info:
|
|
return "amd"
|
|
elif "intel" in gpu_info:
|
|
return "intel"
|
|
except Exception as e:
|
|
logger.warning(f"Could not detect GPU via lspci: {e}")
|
|
|
|
return None
|
|
|
|
def _check_ffmpeg_encoders(self):
|
|
"""Check available encoders in FFmpeg."""
|
|
encoders = {
|
|
"h264_nvenc": False,
|
|
"hevc_nvenc": False,
|
|
"h264_amf": False,
|
|
"hevc_amf": False,
|
|
"h264_qsv": False,
|
|
"hevc_qsv": False,
|
|
"libx264": False,
|
|
"libx265": False,
|
|
}
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["ffmpeg", "-encoders"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5
|
|
)
|
|
output = result.stdout.lower()
|
|
|
|
for encoder in encoders:
|
|
if encoder in output:
|
|
encoders[encoder] = True
|
|
except Exception as e:
|
|
logger.warning(f"Could not check FFmpeg encoders: {e}")
|
|
|
|
return encoders
|
|
|
|
def _recommend_encoder(self):
|
|
"""Recommend best encoder based on hardware."""
|
|
# Prefer NVIDIA NVENC > AMD AMF > Intel QSV > CPU
|
|
|
|
if self.gpu_type == "nvidia":
|
|
if self.available_encoders.get("hevc_nvenc"):
|
|
return "hevc_nvenc" # H.265 on NVIDIA is efficient
|
|
elif self.available_encoders.get("h264_nvenc"):
|
|
return "h264_nvenc"
|
|
|
|
if self.gpu_type == "amd":
|
|
if self.available_encoders.get("hevc_amf"):
|
|
return "hevc_amf"
|
|
elif self.available_encoders.get("h264_amf"):
|
|
return "h264_amf"
|
|
|
|
if self.gpu_type == "intel":
|
|
if self.available_encoders.get("hevc_qsv"):
|
|
return "hevc_qsv"
|
|
elif self.available_encoders.get("h264_qsv"):
|
|
return "h264_qsv"
|
|
|
|
# Fallback to CPU encoders
|
|
if self.available_encoders.get("libx265"):
|
|
return "libx265"
|
|
elif self.available_encoders.get("libx264"):
|
|
return "libx264"
|
|
|
|
return "libx264" # Default fallback
|
|
|
|
def _get_recommended_settings(self):
|
|
"""Get recommended encoding settings based on hardware."""
|
|
settings = {
|
|
"encoder": self.recommended_encoder,
|
|
"preset": self._get_preset(),
|
|
"threads": self._get_thread_count(),
|
|
"gpu_capable": self.gpu_type is not None,
|
|
}
|
|
|
|
return settings
|
|
|
|
def _get_preset(self):
|
|
"""Get recommended preset based on encoder."""
|
|
encoder = self.recommended_encoder
|
|
|
|
# GPU encoders don't use "preset" in the same way
|
|
if "nvenc" in encoder or "amf" in encoder or "qsv" in encoder:
|
|
# GPU presets (0-fast, 1-medium, 2-slow)
|
|
return 1 # medium (balanced)
|
|
else:
|
|
# CPU presets (ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow)
|
|
if self.cpu_cores <= 4:
|
|
return "faster"
|
|
elif self.cpu_cores <= 8:
|
|
return "fast"
|
|
else:
|
|
return "medium"
|
|
|
|
def _get_thread_count(self):
|
|
"""Get recommended thread count for encoding."""
|
|
# For CPU encoding, use most cores but leave some for system
|
|
if "nvenc" in self.recommended_encoder or "amf" in self.recommended_encoder or "qsv" in self.recommended_encoder:
|
|
return 0 # GPU handles it
|
|
else:
|
|
# Leave 1-2 cores for system
|
|
return max(self.cpu_cores - 2, 1)
|
|
|
|
def to_dict(self):
|
|
"""Return all hardware info as dictionary."""
|
|
return {
|
|
"platform": self.platform_name,
|
|
"cpu_cores": self.cpu_cores,
|
|
"gpu_type": self.gpu_type,
|
|
"recommended_encoder": self.recommended_encoder,
|
|
"available_encoders": {k: v for k, v in self.available_encoders.items() if v},
|
|
"recommended_settings": self.recommended_settings,
|
|
}
|
|
|
|
def print_summary(self):
|
|
"""Print hardware detection summary."""
|
|
print("\n" + "="*60)
|
|
print("HARDWARE DETECTION SUMMARY")
|
|
print("="*60)
|
|
print(f"Platform: {self.platform_name}")
|
|
print(f"CPU Cores: {self.cpu_cores}")
|
|
print(f"GPU Type: {self.gpu_type or 'None (CPU only)'}")
|
|
print(f"\nAvailable Encoders:")
|
|
for encoder, available in self.available_encoders.items():
|
|
status = "✓" if available else "✗"
|
|
print(f" {status} {encoder}")
|
|
print(f"\nRecommended:")
|
|
print(f" Encoder: {self.recommended_encoder}")
|
|
print(f" Preset: {self.recommended_settings.get('preset')}")
|
|
print(f" Threads: {self.recommended_settings.get('threads')}")
|
|
print("="*60 + "\n")
|
|
|
|
|
|
def detect_hardware():
|
|
"""Quick detection - returns HardwareInfo object."""
|
|
return HardwareInfo()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Test the hardware detection
|
|
hw = detect_hardware()
|
|
hw.print_summary()
|
|
|
|
# Also save as JSON for reference
|
|
hw_json = hw.to_dict()
|
|
print("Full Hardware Info (JSON):")
|
|
print(json.dumps(hw_json, indent=2))
|