File size: 2,804 Bytes
15d68eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """
ROCm / CUDA health check.
Reports:
- Backend (CUDA vs ROCm)
- GPU count, names, VRAM
- PyTorch version
- Diffusers / Transformers / PEFT versions
- Whether HSA_OVERRIDE_GFX_VERSION is needed (ROCm only)
Resilient: a single library import failure is reported but doesn't crash the check.
"""
from __future__ import annotations
def main():
import torch
print("=" * 60)
print("Indic Heritage Studio v2 — Environment Check")
print("=" * 60)
print(f"\nPyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
backend = "ROCm" if (hasattr(torch.version, "hip")
and torch.version.hip) else "CUDA"
print(f"Backend: {backend}")
if backend == "ROCm":
print(f"ROCm/HIP version: {torch.version.hip}")
print(f"\nGPU count: {torch.cuda.device_count()}")
total_vram = 0.0
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
vram_gb = props.total_memory / 1e9
total_vram += vram_gb
print(f" GPU {i}: {props.name} ({vram_gb:.1f} GB)")
print(f"Total VRAM across all GPUs: {total_vram:.1f} GB")
# Library versions (resilient — one failure shouldn't crash the check)
print("\nLibrary versions:")
for lib in ("diffusers", "transformers", "peft", "accelerate",
"gradio", "safetensors", "controlnet_aux", "cv2",
"compel", "bitsandbytes", "datasets"):
try:
mod = __import__(lib)
ver = getattr(mod, "__version__", "unknown")
print(f" {lib}: {ver}")
except ImportError as e:
print(f" {lib}: NOT INSTALLED ({e})")
except Exception as e:
print(f" {lib}: IMPORT ERROR ({type(e).__name__}: {e})")
# Multi-GPU strategy
if torch.cuda.is_available() and torch.cuda.device_count() >= 4:
print("\n✓ Multi-GPU mode: pipelines will be pinned to dedicated GPUs.")
print(" T2I → GPU 0 | Style → GPU 1 | I2V → GPU 2 | ControlNet → GPU 3")
print(" Batch workers → GPU 4-7")
elif torch.cuda.is_available() and torch.cuda.device_count() >= 2:
print("\n⚠Multi-GPU mode: limited GPUs — pipelines will share.")
else:
print("\n⚠Single-GPU mode — pipelines will load/unload on demand.")
# CUDA_VISIBLE_DEVICES warning
import os
cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
if cvd:
print(f"\n⚠CUDA_VISIBLE_DEVICES is set to '{cvd}' — this restricts GPU visibility.")
print(f" Run: unset CUDA_VISIBLE_DEVICES to see all GPUs.")
print("\n" + "=" * 60)
if __name__ == "__main__":
main()
|