[diffusion] fix: fix temporary resolution workaround (#20046)

python/sglang/multimodal_gen/configs/sample/diffusers_generic.py

@@ -6,7 +6,7 @@ This module provides generic sampling parameters that work with any diffusers pi
 """
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, ClassVar
 
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
@@ -26,6 +26,9 @@ class DiffusersGenericSamplingParams(SamplingParams):
     passed directly to the diffusers pipeline call.
     """
 
+    _default_height: ClassVar[int] = 1024
+    _default_width: ClassVar[int] = 1024
+
     # Override defaults with more conservative values that work across pipelines
     num_frames: int = 1  # default to image generation
     height: int = 1024
@@ -44,9 +47,4 @@ class DiffusersGenericSamplingParams(SamplingParams):
         else:
             self.data_type = DataType.IMAGE
 
-        if self.width is None:
-            self.width_not_provided = True
-            self.width = 1024
-        if self.height is None:
-            self.height_not_provided = True
-            self.height = 1024
+        super().__post_init__()

python/sglang/multimodal_gen/configs/sample/flux.py

@@ -2,25 +2,22 @@
 
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
+from typing import ClassVar
 
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 
 
 @dataclass
 class FluxSamplingParams(SamplingParams):
+    _default_height: ClassVar[int] = 128 * 8  # default_sample_size * vae_scale_factor
+    _default_width: ClassVar[int] = 128 * 8
+
     num_frames: int = 1
     # Denoising stage
     guidance_scale: float = 1.0
     negative_prompt: str = None
     num_inference_steps: int = 50
 
-    def __post_init__(self):
-        default_sample_size = 128
-        vae_scale_factor = 8
-        # FIXME
-        # self.height = default_sample_size * vae_scale_factor
-        # self.width = default_sample_size * vae_scale_factor
-
 
 @dataclass
 class Flux2KleinSamplingParams(FluxSamplingParams):

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -14,7 +14,7 @@ import unicodedata
 import uuid
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields
@@ -123,11 +123,14 @@ class SamplingParams:
     num_frames_round_down: bool = (
         False  # Whether to round down num_frames if it's not divisible by num_gpus
     )
+
+    # Subclasses can set these to provide model-specific default resolutions.
+    # The base __post_init__ will apply them when height/width are not provided.
+    _default_height: ClassVar[int | None] = None
+    _default_width: ClassVar[int | None] = None
+
     height: int | None = None
     width: int | None = None
-    # NOTE: this is temporary, we need a way to know if width or height is not provided, or do the image resize earlier
-    height_not_provided: bool = False
-    width_not_provided: bool = False
     fps: int = 24
 
     # Resolution validation
@@ -217,10 +220,10 @@ class SamplingParams:
     def __post_init__(self) -> None:
         assert self.num_frames >= 1
 
-        if self.width is None:
-            self.width_not_provided = True
-        if self.height is None:
-            self.height_not_provided = True
+        if self.width is None and self._default_width is not None:
+            self.width = self._default_width
+        if self.height is None and self._default_height is not None:
+            self.height = self._default_height
 
         # Handle output_quality to output_compression conversion
         if self.output_compression is None and self.output_quality is not None:
@@ -895,8 +898,6 @@ class SamplingParams:
         sampling_params_fields = {attr.name for attr in dataclasses.fields(cls)}
         args_attrs = set(vars(args).keys())
         attrs = sampling_params_fields & args_attrs
-        args.height_not_provided = False
-        args.width_not_provided = False
         return {attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}
 
     def output_file_path(self):
@@ -927,8 +928,6 @@ class SamplingParams:
                 allow_override_protected or not is_protected_field
             ):
                 setattr(self, field_name, user_value)
-        self.height_not_provided = user_params.height_not_provided
-        self.width_not_provided = user_params.width_not_provided
         self.__post_init__()
 
     @property

python/sglang/multimodal_gen/test/unit/test_sampling_params.py

@@ -1,6 +1,10 @@
 import math
 import unittest
 
+from sglang.multimodal_gen.configs.sample.diffusers_generic import (
+    DiffusersGenericSamplingParams,
+)
+from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 
 
@@ -45,5 +49,23 @@ class TestSamplingParamsValidate(unittest.TestCase):
             SamplingParams(boundary_ratio=math.nan)
 
 
+class TestSamplingParamsSubclass(unittest.TestCase):
+    def test_flux_defaults_resolution_when_not_provided(self):
+        params = FluxSamplingParams()
+
+        self.assertEqual(params.height, 1024)
+        self.assertEqual(params.width, 1024)
+
+    def test_flux_preserves_user_resolution(self):
+        params = FluxSamplingParams(height=640, width=768)
+
+        self.assertEqual(params.height, 640)
+        self.assertEqual(params.width, 768)
+
+    def test_diffusers_generic_calls_base_post_init(self):
+        with self.assertRaises(AssertionError):
+            DiffusersGenericSamplingParams(num_frames=0)
+
+
 if __name__ == "__main__":
     unittest.main()