chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] chore: minor code cleanups (#15190)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Mick
2025-12-15 23:57:02 +08:00
committed by GitHub
parent 7bc8b1532e
commit 1dedb63860
12 changed files with 85 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
python/sglang/multimodal_gen/apps/README.md → python/sglang/multimodal_gen/apps/webui/README.md
View File

6
python/sglang/multimodal_gen/apps/webui/main.py
View File

@@ -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

2
python/sglang/multimodal_gen/docs/profiling.md
View File

@@ -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
```

60
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -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,

49
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -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

56
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -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

15
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -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

3
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -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:
"""

3
python/sglang/multimodal_gen/runtime/models/registry.py
View File

@@ -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

6
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -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

6
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
View File

@@ -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()

6
python/sglang/multimodal_gen/runtime/utils/profiler.py
View File

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