Spaces:
Paused
Paused
File size: 4,545 Bytes
186aa49 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import torch
from ...utils import encode_video, export_to_video
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import InputParam, OutputParam
from .modular_pipeline import Cosmos3OmniModularPipeline
class Cosmos3ActionOutputStep(ModularPipelineBlocks):
model_name = "cosmos3-omni"
@property
def description(self) -> str:
return "Post-processes action latents into action outputs."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
name="action_latents",
type_hint=torch.Tensor,
default=None,
description="Denoised action latents.",
),
InputParam(
name="action_mode", type_hint=str, default=None, description="Requested action-generation mode."
),
InputParam(
name="raw_action_dim_resolved",
type_hint=int,
default=None,
description="Unpadded action-vector dimension.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors.")]
@torch.no_grad()
def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
action_output = None
if block_state.action_mode in {"inverse_dynamics", "policy"} and block_state.action_latents is not None:
action_output = block_state.action_latents
if block_state.raw_action_dim_resolved is not None:
action_output = action_output[:, : block_state.raw_action_dim_resolved]
action_output = [action_output.detach().cpu()]
block_state.action = action_output
self.set_block_state(state, block_state)
return components, state
class Cosmos3ExportStep(ModularPipelineBlocks):
model_name = "cosmos3-omni"
@property
def description(self) -> str:
return (
"Optional export block that writes decoded outputs to disk. Writes `videos` to `output_path` via "
"`export_to_video`, or muxes `videos` with `sound` via `encode_video` when a waveform is present. "
"Not wired into the default blocks; add it explicitly when you want the pipeline to produce a file."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(name="videos", required=True, description="Generated video frames to export."),
InputParam(
name="output_path",
type_hint=str,
required=True,
description="Destination path for the exported video.",
),
InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the exported video."),
InputParam(
name="sound",
type_hint=torch.Tensor,
default=None,
description="Generated waveform to mux into the video.",
),
InputParam(
name="sampling_rate",
type_hint=int,
default=None,
description="Sample rate of the generated waveform in Hz.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("output_path", type_hint=str, description="Path of the exported video file.")]
@torch.no_grad()
def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
output_path = str(block_state.output_path)
fps = int(round(block_state.fps))
if block_state.sound is not None:
if block_state.sampling_rate is None:
raise ValueError("`sampling_rate` is required to export a video with sound.")
encode_video(
block_state.videos,
fps=fps,
audio=block_state.sound,
audio_sample_rate=int(block_state.sampling_rate),
output_path=output_path,
)
else:
export_to_video(block_state.videos, output_path, fps=fps, macro_block_size=1)
block_state.output_path = output_path
self.set_block_state(state, block_state)
return components, state
|