[diffusion] refactor: eliminate redundant parameters in req (#16505)
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user