mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-11-05 14:33:15 +00:00
18 lines
471 B
Python
18 lines
471 B
Python
import os
|
|
from logging import Logger
|
|
|
|
from .logger import TensorboardLogger
|
|
|
|
local_rank = int(os.environ['LOCAL_RANK']) if 'LOCAL_RANK' in os.environ else 0
|
|
world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
|
|
|
|
|
|
def info_if_rank_zero(logger: Logger, msg: str):
|
|
if local_rank == 0:
|
|
logger.info(msg)
|
|
|
|
|
|
def string_if_rank_zero(logger: TensorboardLogger, tag: str, msg: str):
|
|
if local_rank == 0:
|
|
logger.log_string(tag, msg)
|