[diffusion] rl: enable RL rollout path for LTX-2.3 post-training (#28926)

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

@@ -69,6 +69,9 @@ class LTX23SamplingParams(LTX2SamplingParams):
 
     def build_request_extra(self) -> dict[str, Any]:
         extra = super().build_request_extra()
+        # RL rollout uses the official CFG path (guidance_scale=1, no guider).
+        if self.rollout:
+            return extra
         extra["ltx2_stage1_guider_params"] = {
             "video_cfg_scale": self.video_cfg_scale,
             "video_stg_scale": self.video_stg_scale,

python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py

@@ -218,7 +218,9 @@ def _build_response(
 
     responses: list[RolloutResponse] = []
     for sample_idx in range(batch_size):
-        out_i = result.output[sample_idx].contiguous()
+        out_i = result.output[sample_idx]
+        if isinstance(out_i, torch.Tensor):
+            out_i = out_i.contiguous()
         serialized_generated_output = _maybe_serialize(out_i)
         if not rollout:
             responses.append(

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

@@ -5,6 +5,7 @@ from __future__ import annotations
 import base64
 from typing import Any
 
+import numpy as np
 import torch
 from safetensors.torch import load, save
 
@@ -28,6 +29,8 @@ def _maybe_serialize(obj: Any) -> Any:
             "shape": list(obj.shape),
             "dtype": str(obj.dtype),
         }
+    if isinstance(obj, np.ndarray):
+        return _maybe_serialize(torch.from_numpy(obj))
     if isinstance(obj, dict):
         return {k: _maybe_serialize(v) for k, v in obj.items()}
     if isinstance(obj, (list, tuple)):

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -3,7 +3,6 @@ import os
 
 import numpy as np
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
 
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     STAGE_2_DISTILLED_SIGMA_VALUES as _SHARED_STAGE_2_DISTILLED_SIGMA_VALUES,
@@ -23,6 +22,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
     ComponentUse,
     ResidencyState,
 )
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py

@@ -108,6 +108,7 @@ class LTX2AVDecodingStage(DecodingStage):
             trajectory_latents=batch.trajectory_latents,
             trajectory_decoded=None,
             metrics=batch.metrics,
+            rollout_trajectory_data=batch.rollout_trajectory_data,
         )
 
         # 2. Decode Audio

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py

@@ -137,6 +137,12 @@ class LTX2DenoisingStage(DenoisingStage):
         )
         self.sampler_name = sampler_name
 
+    def _scheduler_step_kwargs(self, batch: Req, scheduler) -> dict:
+        return self.prepare_extra_func_kwargs(
+            scheduler.step,
+            {"generator": batch.generator, "eta": batch.eta, "batch": batch},
+        )
+
     @staticmethod
     def _randn_like_with_batch_generators(
         reference_tensor: torch.Tensor, batch: Req
@@ -1873,9 +1879,19 @@ class LTX2DenoisingStage(DenoisingStage):
                     midpoint_model_call=_stage2_midpoint_model_call,
                 )
             else:
-                ctx.latents = ctx.scheduler.step(
+                if batch.rollout:
-                    model_video, step.t_device, ctx.latents, return_dict=False
+                    ctx.scheduler._step_index = step.step_index
-                )[0]
+                    ctx.latents = ctx.scheduler.step(
+                        model_video,
+                        step.t_device,
+                        ctx.latents,
+                        return_dict=False,
+                        **self._scheduler_step_kwargs(batch, ctx.scheduler),
+                    )[0]
+                else:
+                    ctx.latents = ctx.scheduler.step(
+                        model_video, step.t_device, ctx.latents, return_dict=False
+                    )[0]
                 ctx.audio_latents = ctx.audio_scheduler.step(
                     model_audio, step.t_device, ctx.audio_latents, return_dict=False
                 )[0]