From 515d15ef87141198f22db1a2f9494452d0348efe Mon Sep 17 00:00:00 2001 From: Lidang Jiang <119769478+Lidang-Jiang@users.noreply.github.com> Date: Fri, 3 Apr 2026 23:59:07 +0800 Subject: [PATCH] [Cleanup] Replace bare print() with logger and use specific exception types (#2228) Signed-off-by: Lidang Jiang Signed-off-by: Lidang-Jiang Co-authored-by: Claude Opus 4.6 --- .../diffusion/attention/backends/ring/ring_utils.py | 9 ++++++--- vllm_omni/diffusion/diffusion_engine.py | 2 +- vllm_omni/diffusion/models/dreamid_omni/fusion.py | 7 +++++-- .../models/hunyuan_image_3/hunyuan_image_3_tokenizer.py | 5 ++++- .../hunyuan_image_3/hunyuan_image_3_transformer.py | 2 +- vllm_omni/model_executor/models/cosyvoice3/utils.py | 2 +- .../models/qwen3_tts/tokenizer_25hz/vq/core_vq.py | 5 ++++- 7 files changed, 22 insertions(+), 10 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py index c256f62cb..67f71562b 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py @@ -5,6 +5,9 @@ import torch import torch.nn.functional as F +from vllm.logger import init_logger + +logger = init_logger(__name__) __all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"] @@ -79,9 +82,9 @@ def _update_out_and_lse( out = out - F.sigmoid(block_lse - lse) * (out - block_out) lse = lse - F.logsigmoid(lse - block_lse) except RuntimeError as e: - print(f"ERROR in _update_out_and_lse: {e}") - print(f"out: {out.shape}, lse: {lse.shape}") - print(f"block_out: {block_out.shape}, block_lse: {block_lse.shape}") + logger.error("_update_out_and_lse failed: %s", e) + logger.error("out: %s, lse: %s", out.shape, lse.shape) + logger.error("block_out: %s, block_lse: %s", block_out.shape, block_lse.shape) # raise e raise e diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 308c8cef8..05008d7e9 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -115,7 +115,7 @@ class DiffusionEngine: if output.aborted: raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.") if output.error: - raise Exception(f"{output.error}") + raise RuntimeError(f"{output.error}") logger.info("Generation completed successfully.") if output.output is None: diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index 2a4e485fa..a534f5a76 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +from vllm.logger import init_logger from vllm_omni.diffusion.attention.layer import Attention @@ -11,6 +12,8 @@ except ImportError: from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply +logger = init_logger(__name__) + class FusionModel(nn.Module): def __init__(self, video_config=None, audio_config=None): @@ -22,14 +25,14 @@ class FusionModel(nn.Module): else: has_video = False self.video_model = None - print("Warning: No video model is provided!") + logger.warning("No video model is provided!") if audio_config is not None: self.audio_model = WanModel(**audio_config) else: has_audio = False self.audio_model = None - print("Warning: No audio model is provided!") + logger.warning("No audio model is provided!") if has_video and has_audio: assert len(self.video_model.blocks) == len(self.audio_model.blocks) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py index 360904b5e..ce563f711 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py @@ -11,9 +11,12 @@ import torch import torch.nn.functional as F from diffusers.utils.outputs import BaseOutput from transformers import AutoTokenizer +from vllm.logger import init_logger from .hunyuan_image_3_transformer import ImageInfo, JointImageInfo, default +logger = init_logger(__name__) + class TokenizerEncodeOutput(BaseOutput): tokens: torch.Tensor | None = None @@ -121,7 +124,7 @@ class TokenizerWrapper: elif isinstance(uncond_enabled, bool): uncond_enabled = [uncond_enabled] * len(texts) if len(uncond_enabled) != len(texts): - print(uncond_enabled, texts) + logger.debug("uncond_enabled=%s, texts=%s", uncond_enabled, texts) assert len(uncond_enabled) == len(texts), ( f"Length of uncond_flags should be equal to the number of texts, " f"but got {len(uncond_enabled)} and {len(texts)}." diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index 3d670809b..bc81ca9c3 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -2036,7 +2036,7 @@ class HunyuanImage3Model(nn.Module): for name, loaded_weight in weights: # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True) if contains_unexpected_keyword(name, unexpected_keywords): - print(f"Skipping unexpected weight name: {name}") + logger.warning("Skipping unexpected weight name: %s", name) continue if "rotary_emb.inv_freq" in name: continue diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py index e1310cd3b..ca98e9aef 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/utils.py +++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py @@ -180,7 +180,7 @@ def log_mel_spectrogram( HOP_LENGTH = 160 if not torch.is_tensor(audio): - raise Exception(f"audio is not tensor {type(audio)}") + raise TypeError(f"audio is not tensor {type(audio)}") if device is not None: audio = audio.to(device) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py index 9c103a851..5609abb39 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py @@ -40,6 +40,9 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) def round_up_multiple(num, mult): @@ -175,7 +178,7 @@ class EuclideanCodebook(nn.Module): if not torch.any(expired_codes): return else: - print(f"VQ expire infos: num_expire={sum(expired_codes)}, cluster_size[:5]={cluster_size[:5]}") + logger.info("VQ expire infos: num_expire=%s, cluster_size[:5]=%s", sum(expired_codes), cluster_size[:5]) batch_samples = rearrange(batch_samples, "... d -> (...) d") self.replace_(batch_samples, mask=expired_codes)