Wan2GP/shared/default_device.py

from __future__ import annotations

import os
import sys

from collections.abc import MutableSequence


def _arg_name_to_option(arg_name: str) -> str:
    arg_name = str(arg_name or "").strip()
    if not arg_name:
        return ""
    return arg_name if arg_name.startswith("--") else f"--{arg_name}"


def _cuda_visible_device(device: str) -> str:
    device = str(device or "").strip().lower()
    if device.startswith("cuda:"):
        device = device.split(":", 1)[1]
    return device if device.isdigit() else ""


def _rewrite_arg_value(argv: MutableSequence[str], option: str, value: str) -> None:
    for index, arg in enumerate(argv):
        if arg == option and index + 1 < len(argv):
            argv[index + 1] = value
            return
        if str(arg).startswith(f"{option}="):
            argv[index] = f"{option}={value}"
            return


def set_default_cuda_device_from_arg(arg_name: str, default_device: str = "cuda:0") -> bool:
    option = _arg_name_to_option(arg_name)
    if not option:
        return False
    argv = sys.argv
    for index, arg in enumerate(argv[1:], start=1):
        if arg == option and index + 1 < len(argv):
            visible_device = _cuda_visible_device(argv[index + 1])
            break
        if str(arg).startswith(f"{option}="):
            visible_device = _cuda_visible_device(str(arg).split("=", 1)[1])
            break
    else:
        return False

    if not visible_device:
        return False
    os.environ["CUDA_VISIBLE_DEVICES"] = visible_device
    _rewrite_arg_value(argv, option, default_device)
    return True