[diffusion] fix: fix serving with dit-layerwise-offload enabled (#16066)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -152,6 +152,8 @@ class DiffusionServerArgs:
     # LoRA
     lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
 
+    dit_layerwise_offload: bool = False
+
 
 @dataclass(frozen=True)
 class DiffusionSamplingParams:
@@ -171,6 +173,8 @@ class DiffusionSamplingParams:
     # URL direct test flag - if True, don't pre-download URL images
     direct_url_test: bool = False
 
+    num_outputs_per_prompt: int = 1
+
 
 @dataclass(frozen=True)
 class DiffusionTestCase:
@@ -308,6 +312,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    # TODO: replace with a faster model to test the --dit-layerwise-offload
+    # TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice,
+    # so we do one warmup and send one request instead
+    DiffusionTestCase(
+        "flux_2_image_t2i_layerwise_offload",
+        DiffusionServerArgs(
+            model_path="black-forest-labs/FLUX.2-dev",
+            modality="image",
+            dit_layerwise_offload=True,
+            warmup_text=1,
+        ),
+        T2I_sampling_params,
+    ),
     DiffusionTestCase(
         "zimage_image_t2i",
         DiffusionServerArgs(