[diffusion] log: improve logging while multiprocessing (#12997)

This commit is contained in:
Mick
2025-11-12 12:08:37 +08:00
committed by GitHub
parent 151e13687a
commit 60cb716720
4 changed files with 146 additions and 126 deletions
@@ -34,7 +34,10 @@ from huggingface_hub import snapshot_download
from transformers import AutoConfig, PretrainedConfig
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.logging_utils import (
init_logger,
suppress_other_loggers,
)
logger = init_logger(__name__)
_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
@@ -370,7 +373,9 @@ def maybe_download_model(
logger.info(
"Downloading model snapshot from HF Hub for %s...", model_name_or_path
)
with get_lock(model_name_or_path).acquire(poll_interval=2):
with get_lock(model_name_or_path).acquire(
poll_interval=2
), suppress_other_loggers(not_suppress_on_main_rank=True):
local_path = snapshot_download(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],