Skip to content

Commit 0cfc55c

Browse files
yzhautouskayyzhautouskay
committed
Move Cosmos3TransferSetupStep to before_encoder.py
1 parent 2181295 commit 0cfc55c

3 files changed

Lines changed: 167 additions & 161 deletions

File tree

‎src/diffusers/modular_pipelines/cosmos/before_denoise.py‎

Lines changed: 0 additions & 160 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
11
import copy
2-
import math
32

43
import torch
54

6-
from ...configuration_utils import FrozenDict
75
from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
86
from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
97
from ...schedulers import UniPCMultistepScheduler
108
from ...utils.torch_utils import randn_tensor
11-
from ...video_processor import VideoProcessor
129
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
1310
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
1411
from .modular_pipeline import Cosmos3OmniModularPipeline
@@ -965,163 +962,6 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
965962
return components, state
966963

967964

968-
class Cosmos3TransferSetupStep(ModularPipelineBlocks):
969-
model_name = "cosmos3-omni"
970-
971-
@property
972-
def description(self) -> str:
973-
return (
974-
"Preprocesses the transfer control videos and resolves the autoregressive chunk geometry "
975-
"(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop."
976-
)
977-
978-
@property
979-
def expected_components(self) -> list[ComponentSpec]:
980-
return [
981-
ComponentSpec(
982-
"video_processor",
983-
VideoProcessor,
984-
config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
985-
default_creation_method="from_config",
986-
),
987-
]
988-
989-
@property
990-
def inputs(self) -> list[InputParam]:
991-
return [
992-
InputParam(
993-
name="control_videos",
994-
type_hint=dict,
995-
required=True,
996-
description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.",
997-
),
998-
InputParam(
999-
name="height", type_hint=int, default=None, description="Height of the generated video in pixels."
1000-
),
1001-
InputParam(
1002-
name="width", type_hint=int, default=None, description="Width of the generated video in pixels."
1003-
),
1004-
InputParam(
1005-
name="num_frames",
1006-
type_hint=int,
1007-
default=None,
1008-
description="Optional cap on the number of output frames (defaults to the control video length).",
1009-
),
1010-
InputParam(
1011-
name="num_video_frames_per_chunk",
1012-
type_hint=int,
1013-
default=None,
1014-
description="Number of pixel frames generated per autoregressive chunk.",
1015-
),
1016-
InputParam(
1017-
name="num_conditional_frames",
1018-
type_hint=int,
1019-
default=1,
1020-
description="Number of frames each chunk reuses from the previous chunk's tail.",
1021-
),
1022-
]
1023-
1024-
@property
1025-
def intermediate_outputs(self) -> list[OutputParam]:
1026-
return [
1027-
OutputParam("height", type_hint=int, description="Resolved output height in pixels."),
1028-
OutputParam("width", type_hint=int, description="Resolved output width in pixels."),
1029-
OutputParam(
1030-
"control_frames",
1031-
type_hint=dict,
1032-
description="Preprocessed, time-padded control maps in canonical hint order.",
1033-
),
1034-
OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."),
1035-
OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."),
1036-
OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."),
1037-
OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."),
1038-
]
1039-
1040-
@torch.no_grad()
1041-
def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1042-
block_state = self.get_block_state(state)
1043-
device = components._execution_device
1044-
dtype = components.transformer.dtype
1045-
1046-
if block_state.height is None:
1047-
block_state.height = 720
1048-
if block_state.width is None:
1049-
block_state.width = 1280
1050-
1051-
# Canonical hint order used both to validate and to order the preprocessed control maps.
1052-
hint_order = ["edge", "blur", "depth", "seg", "wsm"]
1053-
control_videos = block_state.control_videos
1054-
if not isinstance(control_videos, dict) or not control_videos:
1055-
raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.")
1056-
unknown = [k for k in control_videos if k not in hint_order]
1057-
if unknown:
1058-
raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.")
1059-
if any(v is None for v in control_videos.values()):
1060-
raise ValueError("`control_videos` entries must be loaded videos, not None.")
1061-
1062-
tcf = components.vae_scale_factor_temporal
1063-
sf = components.vae_scale_factor_spatial
1064-
if block_state.height % sf != 0 or block_state.width % sf != 0:
1065-
raise ValueError(
1066-
f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
1067-
)
1068-
1069-
# Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order.
1070-
# The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys).
1071-
hint_keys = [k for k in hint_order if k in control_videos]
1072-
control_frames = {
1073-
key: components.video_processor.preprocess_video(
1074-
control_videos[key], height=block_state.height, width=block_state.width
1075-
).to(device=device, dtype=dtype)
1076-
for key in hint_keys
1077-
}
1078-
1079-
# Output frame count / chunking come from the (first) control video, optionally capped by num_frames.
1080-
total_frames = next(iter(control_frames.values())).shape[2]
1081-
if block_state.num_frames is not None:
1082-
total_frames = min(total_frames, block_state.num_frames)
1083-
total_frames = max(1, total_frames)
1084-
1085-
per_chunk = (
1086-
block_state.num_video_frames_per_chunk
1087-
if block_state.num_video_frames_per_chunk is not None
1088-
else total_frames
1089-
)
1090-
chunk_frames = 1 if total_frames == 1 else per_chunk
1091-
chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1
1092-
1093-
if total_frames <= chunk_frames:
1094-
num_chunks, stride = 1, chunk_frames
1095-
else:
1096-
stride = chunk_frames - block_state.num_conditional_frames
1097-
if stride <= 0:
1098-
raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.")
1099-
remaining = total_frames - chunk_frames
1100-
num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0))
1101-
1102-
# Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to
1103-
# keep reflecting). No truncation here; per-chunk slicing happens later.
1104-
padded = max(total_frames, chunk_frames)
1105-
control_frames_padded = {}
1106-
for key, frames in control_frames.items():
1107-
while frames.shape[2] < padded:
1108-
pad_len = min(frames.shape[2] - 1, padded - frames.shape[2])
1109-
if pad_len <= 0:
1110-
pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1)
1111-
frames = torch.cat([frames, pad_frame], dim=2)
1112-
break
1113-
frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
1114-
control_frames_padded[key] = frames
1115-
block_state.control_frames = control_frames_padded
1116-
block_state.total_frames = total_frames
1117-
block_state.chunk_frames = chunk_frames
1118-
block_state.num_chunks = num_chunks
1119-
block_state.stride = stride
1120-
1121-
self.set_block_state(state, block_state)
1122-
return components, state
1123-
1124-
1125965
class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks):
1126966
model_name = "cosmos3-omni"
1127967

‎src/diffusers/modular_pipelines/cosmos/before_encoder.py‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
import math
2+
3+
import torch
4+
5+
from ...configuration_utils import FrozenDict
6+
from ...video_processor import VideoProcessor
7+
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
8+
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
9+
from .modular_pipeline import Cosmos3OmniModularPipeline
10+
11+
12+
class Cosmos3TransferSetupStep(ModularPipelineBlocks):
13+
model_name = "cosmos3-omni"
14+
15+
@property
16+
def description(self) -> str:
17+
return (
18+
"Preprocesses the transfer control videos and resolves the autoregressive chunk geometry "
19+
"(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop."
20+
)
21+
22+
@property
23+
def expected_components(self) -> list[ComponentSpec]:
24+
return [
25+
ComponentSpec(
26+
"video_processor",
27+
VideoProcessor,
28+
config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
29+
default_creation_method="from_config",
30+
),
31+
]
32+
33+
@property
34+
def inputs(self) -> list[InputParam]:
35+
return [
36+
InputParam(
37+
name="control_videos",
38+
type_hint=dict,
39+
required=True,
40+
description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.",
41+
),
42+
InputParam(
43+
name="height", type_hint=int, default=None, description="Height of the generated video in pixels."
44+
),
45+
InputParam(
46+
name="width", type_hint=int, default=None, description="Width of the generated video in pixels."
47+
),
48+
InputParam(
49+
name="num_frames",
50+
type_hint=int,
51+
default=None,
52+
description="Optional cap on the number of output frames (defaults to the control video length).",
53+
),
54+
InputParam(
55+
name="num_video_frames_per_chunk",
56+
type_hint=int,
57+
default=None,
58+
description="Number of pixel frames generated per autoregressive chunk.",
59+
),
60+
InputParam(
61+
name="num_conditional_frames",
62+
type_hint=int,
63+
default=1,
64+
description="Number of frames each chunk reuses from the previous chunk's tail.",
65+
),
66+
]
67+
68+
@property
69+
def intermediate_outputs(self) -> list[OutputParam]:
70+
return [
71+
OutputParam("height", type_hint=int, description="Resolved output height in pixels."),
72+
OutputParam("width", type_hint=int, description="Resolved output width in pixels."),
73+
OutputParam(
74+
"control_frames",
75+
type_hint=dict,
76+
description="Preprocessed, time-padded control maps in canonical hint order.",
77+
),
78+
OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."),
79+
OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."),
80+
OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."),
81+
OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."),
82+
]
83+
84+
@torch.no_grad()
85+
def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
86+
block_state = self.get_block_state(state)
87+
device = components._execution_device
88+
dtype = components.transformer.dtype
89+
90+
if block_state.height is None:
91+
block_state.height = 720
92+
if block_state.width is None:
93+
block_state.width = 1280
94+
95+
# Canonical hint order used both to validate and to order the preprocessed control maps.
96+
hint_order = ["edge", "blur", "depth", "seg", "wsm"]
97+
control_videos = block_state.control_videos
98+
if not isinstance(control_videos, dict) or not control_videos:
99+
raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.")
100+
unknown = [k for k in control_videos if k not in hint_order]
101+
if unknown:
102+
raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.")
103+
if any(v is None for v in control_videos.values()):
104+
raise ValueError("`control_videos` entries must be loaded videos, not None.")
105+
106+
tcf = components.vae_scale_factor_temporal
107+
sf = components.vae_scale_factor_spatial
108+
if block_state.height % sf != 0 or block_state.width % sf != 0:
109+
raise ValueError(
110+
f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
111+
)
112+
113+
# Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order.
114+
# The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys).
115+
hint_keys = [k for k in hint_order if k in control_videos]
116+
control_frames = {
117+
key: components.video_processor.preprocess_video(
118+
control_videos[key], height=block_state.height, width=block_state.width
119+
).to(device=device, dtype=dtype)
120+
for key in hint_keys
121+
}
122+
123+
# Output frame count / chunking come from the (first) control video, optionally capped by num_frames.
124+
total_frames = next(iter(control_frames.values())).shape[2]
125+
if block_state.num_frames is not None:
126+
total_frames = min(total_frames, block_state.num_frames)
127+
total_frames = max(1, total_frames)
128+
129+
per_chunk = (
130+
block_state.num_video_frames_per_chunk
131+
if block_state.num_video_frames_per_chunk is not None
132+
else total_frames
133+
)
134+
chunk_frames = 1 if total_frames == 1 else per_chunk
135+
chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1
136+
137+
if total_frames <= chunk_frames:
138+
num_chunks, stride = 1, chunk_frames
139+
else:
140+
stride = chunk_frames - block_state.num_conditional_frames
141+
if stride <= 0:
142+
raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.")
143+
remaining = total_frames - chunk_frames
144+
num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0))
145+
146+
# Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to
147+
# keep reflecting). No truncation here; per-chunk slicing happens later.
148+
padded = max(total_frames, chunk_frames)
149+
control_frames_padded = {}
150+
for key, frames in control_frames.items():
151+
while frames.shape[2] < padded:
152+
pad_len = min(frames.shape[2] - 1, padded - frames.shape[2])
153+
if pad_len <= 0:
154+
pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1)
155+
frames = torch.cat([frames, pad_frame], dim=2)
156+
break
157+
frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
158+
control_frames_padded[key] = frames
159+
block_state.control_frames = control_frames_padded
160+
block_state.total_frames = total_frames
161+
block_state.chunk_frames = chunk_frames
162+
block_state.num_chunks = num_chunks
163+
block_state.stride = stride
164+
165+
self.set_block_state(state, block_state)
166+
return components, state

‎src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@
2020
Cosmos3TransferPackSequenceStep,
2121
Cosmos3TransferPrepareLatentsStep,
2222
Cosmos3TransferSetTimestepsStep,
23-
Cosmos3TransferSetupStep,
2423
Cosmos3VisionDenoiseInputStep,
2524
Cosmos3VisionPackSequenceStep,
2625
Cosmos3VisionPrepareLatentsStep,
2726
)
27+
from .before_encoder import Cosmos3TransferSetupStep
2828
from .decoders import (
2929
Cosmos3SoundDecodeStep,
3030
Cosmos3TransferDecodeChunkStep,

0 commit comments

Comments
 (0)