Unverified Commit 58f84489 authored by PengGao's avatar PengGao Committed by GitHub
Browse files

Enhance inference worker to support parallel configuration setup (#556)


Co-authored-by: default avatarhelloyongyang <yongyang1030@163.com>
parent 1d1608b8
import asyncio
import json
import os
from typing import Any, Dict
......@@ -9,7 +8,7 @@ from loguru import logger
from lightx2v.infer import init_runner
from lightx2v.utils.input_info import set_input_info
from lightx2v.utils.set_config import set_config
from lightx2v.utils.set_config import set_config, set_parallel_config
from ..distributed_utils import DistributedManager
......@@ -34,8 +33,12 @@ class TorchrunInferenceWorker:
self.dist_manager.is_initialized = False
config = set_config(args)
if config["parallel"]:
set_parallel_config(config)
if self.rank == 0:
logger.info(f"Config:\n {json.dumps(config, ensure_ascii=False, indent=4)}")
logger.info(f"Config:\n {config}")
self.runner = init_runner(config)
logger.info(f"Rank {self.rank}/{self.world_size - 1} initialization completed")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment