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

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

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

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -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)
+    """
-
+    Create a Req object with sampling_params as a parameter.
-    diffusers_kwargs = params_dict.pop("diffusers_kwargs", None)
+    """
-    extra = params_dict.get("extra") or {}
+    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
+        req.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.adjust_size(server_args)
 

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -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
+    sampling_params: SamplingParams | None = None
-    # properly. Also maybe think about providing an abstraction for pipeline
-    # specific arguments.
-    data_type: DataType
-
-    request_id: str | 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