[REFACT]Refactor diffusion outputs to payload metadata

[bjf-frz]

PLEASE FILL IN THE PR DESCRIPTION HERE.

Purpose

This PR refactors diffusion model outputs from the previous output + custom_output split into a unified payload + metadata envelope.

The goal is to make diffusion outputs have one clear protocol:

• generated data goes to payload
• descriptive information goes to metadata
• temporary model-internal postprocess state can live under metadata["internal"] and is stripped before public formatting
• diffusion runtime code no longer needs a separate custom_output side channel

Previously, DiffusionOutput had two separate output channels:

  DiffusionOutput(
      output=...,
      custom_output=...,
  )

The meaning was not strict.

output usually carried the main generated result:

  output = video_tensor
  output = {"video": video_tensor}
  output = {"image": image_tensor}
  output = {"audio": audio_tensor, "audio_sample_rate": 48000}

custom_output carried extra fields that did not fit into output.

In practice, it included several different kinds of data:

  custom_output = {
      "action": action_tensor,
      "raw_action_dim": 2,
      "action_mode": "policy",
      "domain_id": 7,
      "action_only_output": True,
      "robolab_action_postprocess": ...,
  }

For transfer paths:

  custom_output = {
      "transfer_controls": full_controls,
      "transfer_hints": list(per_hint_frames),
      "fps": frame_rate,
  }

For serving hints:

  custom_output = {
      "video_fps_multiplier": 2,
  }

For some text-like paths:

  custom_output = {
      "text_output": "...",
  }

So the old model was effectively:

  {
      "output": main_output,
      "custom_output": {
          # sometimes generated payload
          # sometimes metadata
          # sometimes serving hints
          # sometimes private postprocess state
          # sometimes model-specific fields
      }
  }

That caused ambiguity because custom_output was not one concept.

Concrete Old Cosmos3 Examples

For Cosmos3 action-only output, the old shape was conceptually:

  DiffusionOutput(
      output={},
      custom_output={
          "action": action,
          "raw_action_dim": raw_action_dim,
          "action_mode": action_mode,
          "domain_id": domain_id,
          "action_only_output": True,
          "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
      },
  )

The problem is that this one dict mixed:

  - actual generated payload: action
  - public metadata: raw_action_dim, action_mode, domain_id
  - serving control metadata: action_only_output
  - internal-only postprocess state: robolab_action_postprocess

For Cosmos3 action + video output, the old shape was conceptually:

  DiffusionOutput(
      output={"video": video},
      custom_output={
          "action": action,
          "raw_action_dim": raw_action_dim,
          "action_mode": action_mode,
          "domain_id": domain_id,
      },
  )

Here video was in output, but action was in custom_output, even though both are generated payloads.

For Cosmos3 transfer output, the old shape was conceptually:

  DiffusionOutput(
      output={"video": full_output},
      custom_output={
          "transfer_controls": full_controls,
          "transfer_hints": list(per_hint_frames),
          "fps": frame_rate,
      },
  )

Here fps is video metadata, but it lived beside transfer controls in custom_output.

———

Why We Should Not Keep custom_output

Keeping custom_output and adding payload / metadata on top would leave two competing output protocols.

That would not really solve the architecture problem.

1. custom_output mixes generated payload and metadata.

Actions are generated outputs. They are closer to image, video, audio, and text than to metadata.

If we kept this:

  output = {"video": video}
  custom_output = {"action": action}

then the runtime still has to know that one generated modality is in output, while another generated modality is in custom_output.

That keeps the original inconsistency.

2. custom_output has no clear ownership boundary.

The old dict could contain:

  "action"
  "raw_action_dim"
  "action_mode"
  "domain_id"
  "action_only_output"
  "robolab_action_postprocess"
  "fps"
  "transfer_controls"
  "transfer_hints"
  "text_output"
  "video_fps_multiplier"

These fields are not the same kind of thing.

Some are payload.
Some are public metadata.
Some are internal metadata.
Some are serving hints.
Some are model-specific temporary state.

A single custom_output dict cannot express those boundaries cleanly.

3. Runtime code would still need model-specific branches.

With custom_output, generic code has to keep doing logic like:

  if "action" in custom_output:
      ...
  if "action_only_output" in custom_output:
      ...
  if "fps" in custom_output:
      ...
  if "text_output" in custom_output:
      ...

That pushes model-specific interpretation into DiffusionEngine, output_formatter, serving_chat, and serving_video.

The refactor goal is the opposite: model-specific details should be normalized by the model pipeline/postprocess, and generic runtime code should consume a stable structure.

4. Postprocess becomes split across two paths.

Before, Cosmos3 had a separate action postprocess hook:

  get_cosmos3_action_post_process_func(...)

That existed because action data lived in custom_output, outside the normal postprocess path.

But action postprocess is still output postprocess. It should be handled inside the model-owned post_process, together with video/image/audio/text formatting.

If custom_output remains, then action still needs special runtime treatment.

5. Internal state cannot be safely separated.

For example:

  "robolab_action_postprocess": ...

is needed to convert model-space actions back to user-facing robot action format.

But it is not public API output.

Under custom_output, this field sits beside public fields like action_mode and domain_id.

With the new structure, it can be represented as:

  "metadata": {
      "internal": {
          "robolab_action_postprocess": ...
      }
  }

Then the formatter strips metadata["internal"] before producing API-facing output.

That gives us a clear private/public boundary.

###After

The new model output shape is:

  DiffusionOutput(
      output={
          "payload": {
              ...
          },
          "metadata": {
              ...
          },
      }
  )

Top-level fields:

  {
      "payload": dict,
      "metadata": dict,
  }

payload contains generated data.

Current supported payload keys include:

  "video"
  "image"
  "audio"
  "text"
  "actions"

Meaning:

  "video": generated video frames or tensor
  "image": generated image output
  "audio": generated audio waveform or audio payload
  "text": generated text output
  "actions": generated action tensor or action payload

metadata contains information needed to interpret payload.

Current metadata groups include:

  "video"
  "audio"
  "actions"
  "transfer"
  "text"
  "common"
  "internal"

Concrete fields:

  "video": {
      "fps": frame_rate,
  }

  "audio": {
      "sample_rate": audio_sample_rate,
  }

  "actions": {
      "raw_action_dim": raw_action_dim,
      "action_mode": action_mode,
      "domain_id": domain_id,
  }

  "transfer": {
      "controls": full_controls,
      "hints": list(per_hint_frames),
  }

  "text": {
      # text-related metadata when needed
  }

  "common": {
      "action_only_output": True,
  }

  "internal": {
      "robolab_action_postprocess": ...
  }

metadata["internal"] is only for model/postprocess-private state. It should not be exposed as public API metadata.

Concrete New Cosmos3 Examples

Action-only output now becomes:

  DiffusionOutput(
      output={
          "payload": {
              "actions": action,
          },
          "metadata": {
              "actions": {
                  "raw_action_dim": raw_action_dim,
                  "action_mode": action_mode,
                  "domain_id": domain_id,
              },
              "common": {
                  "action_only_output": True,
              },
              "internal": {
                  "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
              },
          },
      },
  )

After postprocess, metadata["internal"] is stripped from public output.

Action + video output now becomes:

  DiffusionOutput(
      output={
          "payload": {
              "video": video,
              "actions": action,
          },
          "metadata": {
              "actions": {
                  "raw_action_dim": raw_action_dim,
                  "action_mode": action_mode,
                  "domain_id": domain_id,
              },
          },
      },
  )

Transfer output now becomes:

  DiffusionOutput(
      output={
          "payload": {
              "video": full_output,
          },
          "metadata": {
              "transfer": {
                  "controls": full_controls,
                  "hints": list(per_hint_frames),
              },
              "video": {
                  "fps": frame_rate,
              },
          },
      },
  )

Video + audio output can be represented as:

  {
      "payload": {
          "video": video,
          "audio": audio,
      },
      "metadata": {
          "video": {
              "fps": fps,
          },
          "audio": {
              "sample_rate": sample_rate,
          },
      },
  }

Text output can be represented as:

  {
      "payload": {
          "text": text,
      },
      "metadata": {
          "text": {},
      },
  }

Runtime Behavior After Refactor

DiffusionEngine now only needs one normal postprocess path:

  outputs = self.post_process_func(output_data, **postprocess_kwargs)
  postprocess_output = normalize_diffusion_postprocess_output(outputs)

Then output_formatter normalizes the envelope.

The normalized internal object is:

  DiffusionPostprocessOutput(
      outputs=...,
      metadata=...,
      audio_payload=...,
      action_payload=...,
      audio_sample_rate=...,
      fps=...,
      has_video_payload=...,
  )

Its fields mean:

  outputs:
      primary image/video/text output used by normal OmniRequestOutput formatting

  metadata:
      public metadata after stripping metadata["internal"]

  audio_payload:
      payload["audio"]

  action_payload:
      payload["actions"]

  audio_sample_rate:
      metadata["audio"]["sample_rate"]

  fps:
      metadata["video"]["fps"]

  has_video_payload:
      whether payload contains "video"

This gives generic runtime code one way to consume all diffusion outputs.

Output runtime flow

  Pipeline / Executor
      |
      | returns
      v
  DiffusionOutput
      .output = raw execution result
                torch.Tensor | tuple[Any, ...] | dict[str, Any] | None
      |
      | DiffusionEngine.postprocess_output(...)
      | passes .output into post_process_func
      v
  post_process_func(DiffusionOutput.output)
      |
      | returns
      v
  DiffusionPostprocessRawOutput
      |
      | normalize_diffusion_postprocess_output(...)
      v
  DiffusionPostprocessOutput
      .outputs: DiffusionPayload
      .metadata: DiffusionMetadata
      .primary_key: str | None
      |
      | format_diffusion_outputs(...)
      v
  OmniRequestOutput
      .images
      .final_output_type
      .multimodal_output

Relations between TypeAlias:

1. Metadata Type Graph

DiffusionScalar
DiffusionShape
DiffusionTensorMap
torch.Tensor
list[DiffusionMetadataValue]
tuple[DiffusionMetadataValue, ...]
Mapping[str, DiffusionMetadataValue]
|
v
DiffusionMetadataValue
|
v
DiffusionMetadata
|
v
DiffusionMetadataMapping

Meaning:

DiffusionMetadata
= dict[str, DiffusionMetadataValue]

And:

DiffusionMetadataValue
= DiffusionScalar
| DiffusionShape
| DiffusionTensorMap
| torch.Tensor
| list[DiffusionMetadataValue]
| tuple[DiffusionMetadataValue, ...]
| Mapping[str, DiffusionMetadataValue]

DiffusionMetadataMapping is only the read-only view:

DiffusionMetadataMapping
= Mapping[str, DiffusionMetadataValue]

2. Payload Type Graph

DiffusionScalar
bytes
torch.Tensor
np.ndarray
PIL.Image.Image
list[DiffusionPayloadValue]
tuple[DiffusionPayloadValue, ...]
Mapping[str, DiffusionPayloadValue]
|
v
DiffusionPayloadValue
|
v
DiffusionPayload
|
v
DiffusionPayloadMapping

Meaning:

DiffusionPayload
= dict[str, DiffusionPayloadValue]

And:

DiffusionPayloadValue
= DiffusionScalar
| bytes
| torch.Tensor
| np.ndarray
| PIL.Image.Image
| list[DiffusionPayloadValue]
| tuple[DiffusionPayloadValue, ...]
| Mapping[str, DiffusionPayloadValue]

DiffusionPayloadMapping is the read-only view:

DiffusionPayloadMapping
= Mapping[str, DiffusionPayloadValue]

3. Envelope Type Graph

DiffusionPayload
DiffusionMetadata
|
v
DiffusionOutputEnvelope

Meaning:

DiffusionOutputEnvelope
= {
"payload": DiffusionPayload,
"metadata": DiffusionMetadata,
}

This is the canonical post_process return shape.

4. Raw Output Type Graph

DiffusionPayloadValue
DiffusionPayload
DiffusionOutputEnvelope
|
v
DiffusionPostprocessRawOutput

Meaning:

DiffusionPostprocessRawOutput
= DiffusionPayloadValue
| DiffusionPayload
| DiffusionOutputEnvelope

This means post_process_func may return one of three shapes:

single payload value
OR legacy payload dict
OR canonical envelope

5. Serving Multimodal Output Type Graph

DiffusionPayloadValue
DiffusionMetadata
|
v
DiffusionMultimodalOutput

Meaning:

DiffusionMultimodalOutput
= dict[str, DiffusionPayloadValue | DiffusionMetadata]

This is the type stored in:

OmniRequestOutput.multimodal_output

Minimal Summary Graph

Metadata branch:

DiffusionMetadataValue
|
v
DiffusionMetadata
|
v
DiffusionOutputEnvelope.metadata

Payload branch:

DiffusionPayloadValue
|
v
DiffusionPayload
|
v
DiffusionOutputEnvelope.payload

Accepted postprocess return types:

DiffusionPayloadValue
DiffusionPayload
DiffusionOutputEnvelope
|
v
DiffusionPostprocessRawOutput

Final serving multimodal output:

DiffusionPayloadValue
DiffusionMetadata
|
v
DiffusionMultimodalOutput

Test Plan

vLLM Version:

vLLM-Omni Commit:

Test Result

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[hsliuustc0106]

Please fix pre-commit before review. The current failure is ruff F821: undefined name os in vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:786.

[hsliuustc0106]

Continuing on the proposal/design review: the payload/metadata direction makes sense, but I think two migration issues need to be addressed before full review.

1. Removing action_post_process_func_name from the public out-of-tree register_diffusion_model(...) API needs a migration/deprecation path. Existing plugins that pass this keyword will fail at import time.
2. The custom-output migration is not complete yet. serving_video.py now reads fps/action data only from multimodal_output["metadata"], but existing video-server tests/producers still pass video_fps_multiplier and action metadata through custom_output. Please either migrate those producers/tests to the new envelope or keep a compatibility fallback during the transition.

Also please fill the Test Plan/Test Result with the exact commands and coverage for formatter normalization, Cosmos3 action-only/action+video/transfer, text outputs, serving video/chat, and plugin migration behavior.

[hsliuustc0106]

avoid any datatype, check for all changes

[bjf-frz]

avoid any datatype, check for all changes

Addressed. I tightened the diffusion output protocol typing instead of using broad Any/object at the payload/metadata boundary. The new aliases in output_metadata.py define concrete payload, metadata, envelope, and raw postprocess output types. Remaining dynamic types are kept only at external object boundaries such as engine result objects.

[hsliuustc0106]

why we only add thise function for text output?

[hsliuustc0106]

Can we add a central metadata checklist/validator for the new diffusion output envelope?

Right now the PR defines metadata in prose, but the code mostly treats it as open dict[str, Any]. Since this is becoming the protocol boundary, please add a small canonical checklist for known metadata groups and cover it with unit tests.

Suggested checks:

Group	Key	Check
video	fps	`int
video	video_fps_multiplier	int, >= 1
video	shape	list/tuple of positive ints, if used
audio	sample_rate	int, > 0
audio	audiox_task	non-empty str, if exposed
actions	raw_action_dim	int, > 0
actions	action_mode	non-empty str, ideally one of known action modes
actions	domain_id	int
transfer	controls	dict[str, Any]
transfer	hints	list
text	text_output	str
text	think_text	str
common	action_only_output	bool
internal	any key	allowed only before public formatting; stripped before public output

I would validate known canonical keys without rejecting all unknown model-specific metadata immediately. The important part is preventing typos/shape drift for fields that serving code already consumes, and testing that metadata["internal"] never escapes public formatting.

[hsliuustc0106]

One more requirement for the metadata checklist: please avoid Any in the schema/checklist wherever possible. Using dict[str, Any] or unchecked arbitrary values makes the new protocol hard to understand and easy to misuse, which is the same ambiguity this PR is trying to remove from custom_output.

For canonical metadata fields, the checklist should name concrete value types, for example dict[str, torch.Tensor] for transfer controls if that is the real shape, list[str] for hints, tuple[int, ...] | list[int] for shapes, and explicit scalar types for fps/sample_rate/action fields. Unknown model-specific metadata can still be tolerated at runtime, but the documented/validated fields should not use Any as their contract.

[bjf-frz]

Please fix pre-commit before review. The current failure is ruff F821: undefined name os in vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:786.

Fixed.

[bjf-frz]

Continuing on the proposal/design review: the payload/metadata direction makes sense, but I think two migration issues need to be addressed before full review.

1. Removing action_post_process_func_name from the public out-of-tree register_diffusion_model(...) API needs a migration/deprecation path. Existing plugins that pass this keyword will fail at import time.
2. The custom-output migration is not complete yet. serving_video.py now reads fps/action data only from multimodal_output["metadata"], but existing video-server tests/producers still pass video_fps_multiplier and action metadata through custom_output. Please either migrate those producers/tests to the new envelope or keep a compatibility fallback during the transition.

Also please fill the Test Plan/Test Result with the exact commands and coverage for formatter normalization, Cosmos3 action-only/action+video/transfer, text outputs, serving video/chat, and plugin migration behavior.

Addressed.

• I restored action_post_process_func_name as a deprecated compatibility keyword in register_diffusion_model(...), so out-of-tree plugins that still pass it will not fail at import time. Internally, action handling now goes through the unified pipeline/postprocess output path.
• The remaining diffusion custom-output producers/tests were migrated to the new output envelope. serving_video.py now consumes fps/action-related data from multimodal_output["metadata"], and the producers now emit payload/metadata instead of relying on custom_output.

[hsliuustc0106]

This uses text metadata as the output-type signal, but several image paths now attach text metadata for reasoning/COT (think_text, ar_generated_text) while the payload is still image. Those outputs will be formatted as final_output_type="text" with no images. Please key this off the text payload/primary payload instead.

[hsliuustc0106]

when will we del these trajectory inputs?

[bjf-frz]

when will we del these trajectory inputs?

trajectory_* can be moved into the payload envelope as the canonical path, but we should keep the existing DiffusionOutput.trajectory_* fields as a temporary backward-compatible fallback. New pipelines should write payload["trajectory"], while the formatter can still read legacy trajectory_* fields until downstream custom pipelines/tests are migrated.

[hsliuustc0106]

add the functions for this file and all other new files from line 3

[bjf-frz]

add the functions for this file and all other new files from line 3

done

[hsliuustc0106]

fix ci, cc @SamitHuang @wtomin verl-omni ci breaks

[hsliuustc0106]

lgtm