video-dubbing-agent / utils /gpu_detector.py
dashhdata's picture
Upload folder using huggingface_hub
4ec3855 verified
"""
Utility — GPU Detection
Auto-detect CUDA availability and recommend settings.
"""
import logging
import subprocess
logger = logging.getLogger(__name__)
def detect_gpu() -> dict:
"""
Detect GPU availability and return system info.
"""
info = {
"cuda_available": False,
"gpu_name": None,
"gpu_memory_mb": 0,
"recommended_device": "cpu",
"recommended_compute": "int8",
"recommended_batch_size": 4,
}
try:
import torch
if torch.cuda.is_available():
info["cuda_available"] = True
info["gpu_name"] = torch.cuda.get_device_name(0)
info["gpu_memory_mb"] = torch.cuda.get_device_properties(0).total_mem // (1024 * 1024)
info["recommended_device"] = "cuda"
info["recommended_compute"] = "float16"
info["recommended_batch_size"] = 16 if info["gpu_memory_mb"] > 8000 else 8
logger.info(f"GPU detected: {info['gpu_name']} ({info['gpu_memory_mb']} MB)")
else:
logger.info("PyTorch installed but no CUDA GPU detected.")
except ImportError:
logger.info("PyTorch not installed. Using CPU mode.")
# Fallback: check nvidia-smi
if not info["cuda_available"]:
try:
result = subprocess.run(
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
capture_output=True, text=True, timeout=5
)
if result.returncode == 0 and result.stdout.strip():
parts = result.stdout.strip().split(",")
info["gpu_name"] = parts[0].strip()
logger.info(f"nvidia-smi found GPU: {info['gpu_name']} (but PyTorch CUDA not available)")
except (FileNotFoundError, subprocess.TimeoutExpired):
pass
return info
def get_system_report() -> str:
"""Generate a human-readable system readiness report."""
gpu = detect_gpu()
lines = [
"=== System Report ===",
f"GPU: {gpu['gpu_name'] or 'None detected'}",
f"CUDA: {'Yes' if gpu['cuda_available'] else 'No'}",
f"VRAM: {gpu['gpu_memory_mb']} MB" if gpu['gpu_memory_mb'] else "VRAM: N/A",
f"Mode: {gpu['recommended_device'].upper()}",
f"Compute: {gpu['recommended_compute']}",
f"Batch size: {gpu['recommended_batch_size']}",
"",
"Estimated processing times for 2.5hr video:",
]
if gpu["cuda_available"] and gpu["gpu_memory_mb"] > 8000:
lines.append(" Transcription: ~15-25 min")
lines.append(" TTS Generation: ~20-30 min")
lines.append(" Total: ~45-75 min")
elif gpu["cuda_available"]:
lines.append(" Transcription: ~25-40 min")
lines.append(" TTS Generation: ~30-45 min")
lines.append(" Total: ~75-120 min")
else:
lines.append(" Transcription: ~2-4 hours (CPU)")
lines.append(" TTS Generation: ~1-2 hours (CPU)")
lines.append(" Total: ~3-6 hours")
lines.append(" ⚠ GPU strongly recommended for videos > 30 min")
return "\n".join(lines)