[Cleanup] Replace bare print() with logger and use specific exception types (#2228)

Signed-off-by: Lidang Jiang <lidangjiang@gmail.com>
Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

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
 

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:

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)

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)}."

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

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)

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)