[Core][Diffusion] Add a typed pre-D2H video media contract (#6615)

Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
Co-authored-by: Samit <285365963@qq.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -47,6 +47,7 @@ from vllm_omni.diffusion.output_formatter import (
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
 )
+from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
     get_diffusion_post_process_func,
@@ -534,26 +535,33 @@ class DiffusionEngine:
             raise RuntimeError(output.error)
         logger.debug("Generation completed successfully.")
 
-        if output.output is None:
-            logger.warning("Output is None, returning empty OmniRequestOutput")
-            return format_empty_diffusion_outputs(request, finished=output.finished)
-
-        # When CPU offload is enabled, move output to CPU before
-        # post-processing to avoid device OOM — model weights may still
-        # reside on the device and leave no headroom for intermediates.
-        output_data = output.output
-        if self.od_config.enable_cpu_offload:
-            output_data = _move_tensor_tree_to_cpu(output_data)
-
-        if self.post_process_func is not None:
-            # Some video pipelines need request-level controls during
-            # postprocess (for example worker-side frame interpolation).
-            postprocess_kwargs: dict[str, object] = {}
-            if self._post_process_accepts_sampling_params:
-                postprocess_kwargs["sampling_params"] = request.sampling_params
-            outputs = self.post_process_func(output_data, **postprocess_kwargs)
+        if output.media is not None:
+            if output.output is not None:
+                raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+            media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media
+            output_data = media.video.tensor
+            outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params)
         else:
-            outputs = output_data
+            if output.output is None:
+                logger.warning("Output is None, returning empty OmniRequestOutput")
+                return format_empty_diffusion_outputs(request, finished=output.finished)
+
+            # When CPU offload is enabled, move output to CPU before
+            # post-processing to avoid device OOM — model weights may still
+            # reside on the device and leave no headroom for intermediates.
+            output_data = output.output
+            if self.od_config.enable_cpu_offload:
+                output_data = _move_tensor_tree_to_cpu(output_data)
+
+            if self.post_process_func is not None:
+                # Some video pipelines need request-level controls during
+                # postprocess (for example worker-side frame interpolation).
+                postprocess_kwargs: dict[str, object] = {}
+                if self._post_process_accepts_sampling_params:
+                    postprocess_kwargs["sampling_params"] = request.sampling_params
+                outputs = self.post_process_func(output_data, **postprocess_kwargs)
+            else:
+                outputs = output_data
 
         postprocess_output = normalize_diffusion_postprocess_output(outputs)