[diffusion] refactor: eliminate redundant parameters in req (#16505)

Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Yuhao Yang
2026-01-08 11:14:03 +08:00
committed by GitHub
parent 8867d24879
commit e14f5ec8a8
3 changed files with 60 additions and 87 deletions

View File

@@ -1,7 +1,6 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import base64
import dataclasses
import os
import time
from typing import List, Optional
@@ -27,11 +26,9 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
save_image_to_path,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import shallow_asdict
router = APIRouter(prefix="/v1/images", tags=["images"])
logger = init_logger(__name__)
@@ -104,11 +101,6 @@ def _build_sampling_params_from_request(
return sampling_params
def _build_req_from_sampling(s: SamplingParams) -> Req:
req_fields = {f.name for f in dataclasses.fields(Req)}
return Req(**{k: v for k, v in shallow_asdict(s).items() if k in req_fields})
@router.post("/generations", response_model=ImageResponse)
async def generations(
request: ImageGenerationsRequest,
@@ -245,7 +237,10 @@ async def edits(
enable_teacache=enable_teacache,
num_frames=num_frames,
)
batch = _build_req_from_sampling(sampling)
batch = prepare_request(
server_args=get_global_server_args(),
sampling_params=sampling,
)
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch

View File

@@ -8,7 +8,6 @@ This module provides a consolidated interface for generating videos using
diffusion models.
"""
import dataclasses
import os
import imageio
@@ -21,7 +20,6 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
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 CYAN, RESET, init_logger
from sglang.multimodal_gen.utils import shallow_asdict
logger = init_logger(__name__)
@@ -30,22 +28,13 @@ def prepare_request(
server_args: ServerArgs,
sampling_params: SamplingParams,
) -> Req:
params_dict = shallow_asdict(sampling_params)
diffusers_kwargs = params_dict.pop("diffusers_kwargs", None)
extra = params_dict.get("extra") or {}
"""
Create a Req object with sampling_params as a parameter.
"""
req = Req(sampling_params=sampling_params, VSA_sparsity=server_args.VSA_sparsity)
diffusers_kwargs = getattr(sampling_params, "diffusers_kwargs", None)
if diffusers_kwargs:
extra["diffusers_kwargs"] = diffusers_kwargs
params_dict["extra"] = extra
# Filter to only fields that exist in Req
req_fields = {f.name for f in dataclasses.fields(Req)}
filtered_params = {k: v for k, v in params_dict.items() if k in req_fields}
req = Req(
**filtered_params,
VSA_sparsity=server_args.VSA_sparsity,
)
req.extra["diffusers_kwargs"] = diffusers_kwargs
req.adjust_size(server_args)

View File

@@ -19,7 +19,7 @@ from typing import Any, Optional
import PIL.Image
import torch
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
from sglang.multimodal_gen.configs.sample.teacache import (
TeaCacheParams,
WanTeaCacheParams,
@@ -40,19 +40,15 @@ class Req:
This dataclass contains all information needed during the diffusion pipeline
execution, allowing methods to update specific components without needing
to manage numerous individual parameters.
[IMPORTANT] Fields that overlap with SamplingParams are automatically delegated to the
sampling_params member via __getattr__ and __setattr__.
"""
# TODO(will): double check that args are separate from server_args
# properly. Also maybe think about providing an abstraction for pipeline
# specific arguments.
data_type: DataType
request_id: str | None = None
sampling_params: SamplingParams | None = None
generator: torch.Generator | list[torch.Generator] | None = None
# Image inputs
image_path: str | list[str] | None = None
# Image encoder hidden states
image_embeds: list[torch.Tensor] = field(default_factory=list)
@@ -62,13 +58,6 @@ class Req:
pixel_values: torch.Tensor | PIL.Image.Image | None = None
preprocessed_image: torch.Tensor | None = None
# Text inputs
prompt: str | list[str] | None = None
negative_prompt: str | list[str] | None = None
prompt_path: str | None = None
output_path: str = "outputs/"
# without extension
output_file_name: str | None = None
output_file_ext: str | None = None
# Primary encoder embeddings
prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list)
@@ -86,13 +75,7 @@ class Req:
prompt_template: dict[str, Any] | None = None
do_classifier_free_guidance: bool = False
# Batch info
num_outputs_per_prompt: int = 1
seed: int | None = 42
seeds: list[int] | None = None
generator_device: str = (
"cuda" # Device for random generator: "cuda", "musa" or "cpu"
)
# Tracking if embeddings are already processed
is_prompt_processed: bool = False
@@ -111,35 +94,12 @@ class Req:
# Latent dimensions
height_latents: list[int] | int | None = None
width_latents: list[int] | int | None = None
num_frames: list[int] | int = 1 # Default for image models
num_frames_round_down: bool = (
False # Whether to round down num_frames if it's not divisible by num_gpus
)
# Original dimensions (before VAE scaling)
height: list[int] | int | None = None
width: list[int] | int | None = None
fps: list[int] | int | None = None
height_not_provided: bool = False
width_not_provided: bool = False
# Timesteps
timesteps: torch.Tensor | None = None
timestep: torch.Tensor | float | int | None = None
step_index: int | None = None
boundary_ratio: float | None = None
# Scheduler parameters
# Can be overridden via SGLANG_TEST_NUM_INFERENCE_STEPS env var for faster testing
num_inference_steps: int = int(
os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS", "50")
)
guidance_scale: float = 1.0
guidance_scale_2: float | None = None
true_cfg_scale: float | None = (
None # for CFG vs guidance distillation (e.g., QwenImage)
)
guidance_rescale: float = 0.0
eta: float = 0.0
sigmas: list[float] | None = None
@@ -151,21 +111,15 @@ class Req:
# Component modules (populated by the pipeline)
modules: dict[str, Any] = field(default_factory=dict)
return_trajectory_latents: bool = False
return_trajectory_decoded: bool = False
trajectory_timesteps: list[torch.Tensor] | None = None
trajectory_latents: torch.Tensor | None = None
# Extra parameters that might be needed by specific pipeline implementations
extra: dict[str, Any] = field(default_factory=dict)
# Misc
save_output: bool = True
return_frames: bool = False
is_warmup: bool = False
# TeaCache parameters
enable_teacache: bool = False
teacache_params: TeaCacheParams | WanTeaCacheParams | None = None
# STA parameters
@@ -180,19 +134,54 @@ class Req:
# stage logging
timings: Optional["RequestTimings"] = None
# profile
profile: bool = False
profile_all_stages: bool = False
num_profiled_timesteps: int = None
# debugging
debug: bool = False
# dummy for now
perf_dump_path: str | None = None
# results
output: torch.Tensor | None = None
def __getattr__(self, name: str) -> Any:
"""
Delegate attribute access to sampling_params if not found in Req.
This is only called when the attribute is not found in the instance.
"""
if name == "sampling_params":
raise AttributeError(
f"'{type(self).__name__}' object has no attribute '{name}'"
)
sampling_params = object.__getattribute__(self, "sampling_params")
if sampling_params is not None and hasattr(sampling_params, name):
return getattr(sampling_params, name)
raise AttributeError(
f"'{type(self).__name__}' object has no attribute '{name}'"
)
def __setattr__(self, name: str, value: Any) -> None:
"""
Smart attribute setting:
1. If field exists in Req, set it in Req
2. Else if field exists in sampling_params, set it in sampling_params
3. Else set it in Req (for dynamic attributes)
"""
if name == "sampling_params":
object.__setattr__(self, name, value)
return
if name in self.__class__.__dataclass_fields__:
object.__setattr__(self, name, value)
return
try:
sampling_params = object.__getattribute__(self, "sampling_params")
if sampling_params is not None and hasattr(sampling_params, name):
setattr(sampling_params, name, value)
return
except AttributeError:
# This can happen if `sampling_params` is not set yet. We'll fall through
# to setting the attribute on `self`.
pass
object.__setattr__(self, name, value)
@property
def batch_size(self):
# Determine batch size