From 1dedb63860691a6ae01d0334597d2c2a0fccf2ad Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 15 Dec 2025 23:57:02 +0800 Subject: [PATCH] [diffusion] chore: minor code cleanups (#15190) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../multimodal_gen/apps/{ => webui}/README.md | 0 .../sglang/multimodal_gen/apps/webui/main.py | 6 +- .../sglang/multimodal_gen/docs/profiling.md | 2 +- .../entrypoints/diffusion_generator.py | 60 ++----------------- .../runtime/entrypoints/openai/utils.py | 49 +-------------- .../runtime/entrypoints/utils.py | 56 ++++++++++++++++- .../runtime/loader/component_loader.py | 15 +++-- .../runtime/managers/gpu_worker.py | 3 - .../multimodal_gen/runtime/models/registry.py | 3 +- .../pipelines_core/stages/denoising.py | 6 +- .../runtime/utils/logging_utils.py | 6 +- .../multimodal_gen/runtime/utils/profiler.py | 6 +- 12 files changed, 85 insertions(+), 127 deletions(-) rename python/sglang/multimodal_gen/apps/{ => webui}/README.md (100%) diff --git a/python/sglang/multimodal_gen/apps/README.md b/python/sglang/multimodal_gen/apps/webui/README.md similarity index 100% rename from python/sglang/multimodal_gen/apps/README.md rename to python/sglang/multimodal_gen/apps/webui/README.md diff --git a/python/sglang/multimodal_gen/apps/webui/main.py b/python/sglang/multimodal_gen/apps/webui/main.py index 2257adac9..be60dcfb9 100644 --- a/python/sglang/multimodal_gen/apps/webui/main.py +++ b/python/sglang/multimodal_gen/apps/webui/main.py @@ -5,8 +5,10 @@ from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) -from sglang.multimodal_gen.runtime.entrypoints.openai.utils import post_process_sample -from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + post_process_sample, + prepare_request, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client diff --git a/python/sglang/multimodal_gen/docs/profiling.md b/python/sglang/multimodal_gen/docs/profiling.md index b675816da..3b8c1342e 100644 --- a/python/sglang/multimodal_gen/docs/profiling.md +++ b/python/sglang/multimodal_gen/docs/profiling.md @@ -45,7 +45,7 @@ sglang generate \ By default, trace files are saved in the ./logs/ directory. The exact output file path will be shown in the console output, for example: ```bash -[mm-dd hh:mm:ss] Saving profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz +[mm-dd hh:mm:ss] Saved profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz ``` {request_id}-{num_steps}_steps-global-rank{rank}.trace.json.gz ``` diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 83b4409ee..f1075d190 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -13,22 +13,18 @@ import os import time from typing import Any -import imageio import numpy as np -import torch -import torchvision -from einops import rearrange -from sglang.multimodal_gen.configs.sample.sampling_params import ( - DataType, - SamplingParams, -) +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( MergeLoraWeightsReq, SetLoraReq, UnmergeLoraWeightsReq, ) -from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + post_process_sample, + prepare_request, +) from sglang.multimodal_gen.runtime.launch_server import launch_server from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch @@ -39,7 +35,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import ( log_batch_completion, log_generation_timer, suppress_loggers, - suppress_other_loggers, ) suppress_loggers(["imageio", "imageio_ffmpeg", "PIL", "PIL_Image"]) @@ -162,49 +157,6 @@ class DiffGenerator: f"{self.server_args.scheduler_endpoint()}." ) - def post_process_sample( - self, - sample: torch.Tensor, - data_type: DataType, - fps: int, - save_output: bool = True, - save_file_path: str = None, - ): - """ - Process a single sample output and save output if necessary - """ - # Process outputs - if sample.dim() == 3: - # for images, dim t is missing - sample = sample.unsqueeze(1) - sample = rearrange(sample, "c t h w -> t c h w") - frames = [] - # TODO: this can be batched - for x in sample: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save outputs if requested - if save_output: - if save_file_path: - os.makedirs(os.path.dirname(save_file_path), exist_ok=True) - with suppress_other_loggers(): - if data_type == DataType.VIDEO: - imageio.mimsave( - save_file_path, - frames, - fps=fps, - format=data_type.get_default_extension(), - ) - else: - imageio.imwrite(save_file_path, frames[0]) - logger.info("Saved output to %s", save_file_path) - else: - logger.warning("No output path provided, output not saved") - - return frames - def generate( self, sampling_params_kwargs: dict | None = None, @@ -280,7 +232,7 @@ class DiffGenerator: continue for output_idx, sample in enumerate(output_batch.output): num_outputs = len(output_batch.output) - frames = self.post_process_sample( + frames = post_process_sample( sample, fps=req.fps, save_output=req.save_output, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 4e09b34df..01a54e05f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -4,14 +4,9 @@ import os import time from typing import Optional -import imageio -import numpy as np -import torch -import torchvision -from einops import rearrange from fastapi import UploadFile -from sglang.multimodal_gen.configs.sample.sampling_params import DataType +from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, log_batch_completion, @@ -38,47 +33,6 @@ class UnmergeLoraWeightsReq: target: str = "all" # "all", "transformer", "transformer_2", "critic" -def post_process_sample( - sample: torch.Tensor, - data_type: DataType, - fps: int, - save_output: bool = True, - save_file_path: str = None, -): - """ - Process sample output and save video if necessary - """ - # Process outputs - if sample.dim() == 3: - # for images, dim t is missing - sample = sample.unsqueeze(1) - videos = rearrange(sample, "c t h w -> t c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save outputs if requested - if save_output: - if save_file_path: - os.makedirs(os.path.dirname(save_file_path), exist_ok=True) - if data_type == DataType.VIDEO: - imageio.mimsave( - save_file_path, - frames, - fps=fps, - format=data_type.get_default_extension(), - ) - else: - imageio.imwrite(save_file_path, frames[0]) - logger.info(f"Saved output to {save_file_path}") - else: - logger.info(f"No output path provided, output not saved") - - return frames - - def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]: try: parts = size.lower().replace(" ", "").split("x") @@ -87,7 +41,6 @@ def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]: w, h = int(parts[0]), int(parts[1]) return w, h except Exception: - # Fallback to default portrait 720x1280 return None, None diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 7d84602a9..2d2f2f9ca 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -9,11 +9,21 @@ diffusion models. """ import dataclasses +import os -from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger from sglang.multimodal_gen.utils import shallow_asdict logger = init_logger(__name__) @@ -44,3 +54,45 @@ def prepare_request( ) return req + + +def post_process_sample( + sample: torch.Tensor, + data_type: DataType, + fps: int, + save_output: bool = True, + save_file_path: str = None, +): + """ + Process sample output and save video if necessary + """ + # Process outputs + if sample.dim() == 3: + # for images, dim t is missing + sample = sample.unsqueeze(1) + videos = rearrange(sample, "c t h w -> t c h w") + frames = [] + # TODO: this can be batched + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + # Save outputs if requested + if save_output: + if save_file_path: + os.makedirs(os.path.dirname(save_file_path), exist_ok=True) + if data_type == DataType.VIDEO: + imageio.mimsave( + save_file_path, + frames, + fps=fps, + format=data_type.get_default_extension(), + ) + else: + imageio.imwrite(save_file_path, frames[0]) + logger.info(f"Saved output to {CYAN}{save_file_path}{RESET}") + else: + logger.info(f"No output path provided, output not saved") + + return frames diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index db71023a6..4fdc9cdf1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -128,11 +128,16 @@ class ComponentLoader(ABC): component_model_path, server_args, module_name ) source = "customized" - except Exception as _e: - traceback.print_exc() - logger.error( - f"Error while loading customized {module_name}, falling back to native version" - ) + except Exception as e: + if "Unsupported model architecture" in str(e): + logger.info( + f"Module: {module_name} doesn't have a customized version yet, using native version" + ) + else: + traceback.print_exc() + logger.error( + f"Error while loading customized {module_name}, falling back to native version" + ) # fallback to native version component = self.load_native( component_model_path, server_args, transformers_or_diffusers diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 0482af743..b9bab94fa 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -33,9 +33,6 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import ( logger = init_logger(__name__) -CYAN = "\033[1;36m" -RESET = "\033[0;0m" - class GPUWorker: """ diff --git a/python/sglang/multimodal_gen/runtime/models/registry.py b/python/sglang/multimodal_gen/runtime/models/registry.py index 2ee262571..5e6367a40 100644 --- a/python/sglang/multimodal_gen/runtime/models/registry.py +++ b/python/sglang/multimodal_gen/runtime/models/registry.py @@ -316,8 +316,9 @@ class _ModelRegistry: normalized_arch = [] for arch in architectures: if arch not in self.registered_models: + registered_models = list(self.registered_models.keys()) raise Exception( - f"Unsupported model architecture: {arch}. Registered architectures: {self.registered_models=}" + f"Unsupported model architecture: {arch}. Registered architectures: {registered_models}" ) normalized_arch.append(arch) return normalized_arch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 81af951c3..8f30f5aeb 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1114,10 +1114,8 @@ class DenoisingStage(PipelineStage): A tqdm progress bar. """ local_rank = get_world_group().local_rank - if local_rank == 0: - return tqdm(iterable=iterable, total=total) - else: - return tqdm(iterable=iterable, total=total, disable=True) + disable = local_rank != 0 + return tqdm(iterable=iterable, total=total, disable=disable) def rescale_noise_cfg( self, noise_cfg, noise_pred_text, guidance_rescale=0.0 diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index a4b1393bc..9e55b72f0 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -24,6 +24,8 @@ SGLANG_DIFFUSION_LOGGING_CONFIG_PATH = envs.SGLANG_DIFFUSION_LOGGING_CONFIG_PATH SGLANG_DIFFUSION_LOGGING_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX +# color +CYAN = "\033[1;36m" RED = "\033[91m" GREEN = "\033[92m" YELLOW = "\033[93m" @@ -484,10 +486,6 @@ def log_generation_timer( total_requests, prompt[:100], ) - else: - max_len = 100 - suffix = "..." if len(prompt) > max_len else "" - logger.info(f"Processing prompt: {prompt[:100]}{suffix}") timer = GenerationTimer() timer.start_time = time.perf_counter() diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index 2bf9c283e..18f8ad553 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/python/sglang/multimodal_gen/runtime/utils/profiler.py @@ -2,7 +2,7 @@ import os import torch -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger logger = init_logger(__name__) @@ -127,7 +127,7 @@ class SGLDiffusionProfiler: f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{dump_rank}.trace.json.gz", ) ) - logger.info(f"Saving profiler traces to: {trace_path}") self.profiler.export_chrome_trace(trace_path) + logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}") except Exception as e: - logger.error(f"Failed to export trace: {e}") + logger.error(f"Failed to save trace: {e}")