Cosmos3 edge support#14181
Conversation
github-actions
Bot
added
models
tests
modular-pipelines
pipelines
size/L
|
Hi @atharvajoshi10, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@atharvajoshi10 can you look into the test failing? I think the repo was removed |
yiyixuxu
left a comment
yiyixuxu
left a comment
There was a problem hiding this comment.
my bad I completely forgot review the tests on last PR so I reviewed the entire file, left some comments on tests added from previous PR as well
| assert state.get("uncond_packed_static") is None | ||
|
|
||
|
|
||
| def test_cosmos3_text_step_uses_pipeline_system_prompt_and_safety_configs(monkeypatch): |
There was a problem hiding this comment.
can you just move this test to the ModularPipelineTesterMixin?
if you want to best the behavior of use_system_prompt you can do something like this
text_encoder_block = Cosmos3TextEncoderStep()
text_pipe = text_encoder_block.init_pipeline(self.pretrained_model_name_or_path)
text_pipe.load_components()
text_pipe.disable_safety_checker()
assert text_pipe.config_default_use_system_prompt == True
ids1 = text_pipe(prompt=..., output="cond_input_ids")
ids2 = text_pipe(prompt = , use_system_prompt=True, output="cond_input_ids")
ids3 = text_pipe(prompt = , use_system_prompt=False, output="cond_input_ids")
text_pipe.update_components(
default_use_system_prompt=default_use_system_prompt,
)
assert text_pipe.config_default_use_system_prompt == False
ids4= text_pipe(prompt=... output="cond_input_ids")
ids5 = text_pipe(prompt = , use_system_prompt=True, output="cond_input_ids")
ids6 = text_pipe(prompt = , use_system_prompt=False, output="cond_input_ids")
# compare
assert list(ids1) == list(ids2) == list(ids5) # ids with system prompt
assert list(ids3) == list(ids4) == list(ids6) # ids without system prompt
assert len(ids1) > len(id3) # system prompt adds tokens| assert state.get("cond_input_ids").shape == (1, 1) | ||
|
|
||
|
|
||
| def test_cosmos3_native_flow_schedule_uses_edge_sigma_grid(monkeypatch): |
There was a problem hiding this comment.
similarly, can we move this test to TestCosmos3OmniModularPipelineFast as well?
def test_set_timesteps_native_flow_schedule(self):
timesteps_pipe = Cosmos3SetTimestepsStep().init_pipeline(self.pretrained_model_name_or_path)
timesteps_pipe.load_components()
assert timesteps_pipe.config.use_native_flow_schedule is False
timesteps_default = timesteps_pipe(num_inference_steps=4, output="timesteps")
timesteps_pipe.update_components(
scheduler=UniPCMultistepScheduler(num_train_timesteps=100, use_flow_sigmas=True),
use_native_flow_schedule=True,
)
timesteps_native = timesteps_pipe(num_inference_steps=4, output="timesteps")
expected_sigmas = ...
torch.testing.assert_close(
timesteps_pipe.scheduler.sigmas[:-1], expected_sigmas
)
assert timesteps_native.tolist() == [99, 74, 49, 24]
assert not torch.equal(timesteps_native, timesteps_default)| pipe(**inputs, output=self.output_name) | ||
|
|
||
|
|
||
| def test_cosmos3_pack_steps_do_not_require_vae(): |
There was a problem hiding this comment.
can you move this test to the mixin too?
you can create your packing block easier
pack_block = Cosmos3VisionPackSequenceStep()
pack_pipe = pack_block.init_pipeline(self.pretrained_model_name_or_path)
pack_pipe.load_components()| assert action_segment["action_mrope_ids"].shape == (3, 4) | ||
|
|
||
|
|
||
| def test_cosmos3_denoise_input_steps_assemble_modality_segments(): |
There was a problem hiding this comment.
can you move this to the mixin too
and use the correct API to run them, i.e. pipeline blocks are not meant to be executed, even though you can if you directly pass it a state; the API to run a block is to turn it into a pipeline and then you can pass input directly, it is more meaningful to test the execution API
blocks = SequentialPipelineBlocks.from_blocks_dict(
{
"vision": Cosmos3VisionDenoiseInputStep(),
"sound": Cosmos3SoundDenoiseInputStep(),
"action": Cosmos3ActionDenoiseInputStep(),
}
)
pipe = blocks.init_pipeline() # no repo needed — these blocks declare no components
state = pipe(
cond_text_segment=cond_text_segment,
uncond_text_segment=uncond_text_segment,
... # all eight segments, passed as normal call kwargs
)
torch.testing.assert_close(state.get("cond_position_ids"), torch.cat([...], dim=1))
assert state.get("cond_sequence_length") == 6| assert state.get("timesteps").tolist() == [0, 1, 2, 3] | ||
|
|
||
|
|
||
| def test_cosmos3_modular_model_index_takes_precedence(tmp_path): |
There was a problem hiding this comment.
this is testing loading behavior of base pipeline, can you remove it?
we can add a test to our base test to test it
| @@ -338,9 +394,22 @@ | |||
|
|
|||
|
|
|||
| def test_cosmos3_model_index_fallback_resolves_modular_pipeline(tmp_path): | |||
There was a problem hiding this comment.
can you remove this too?
| state.set("num_frames", 5) | ||
| state.set("height", 32) | ||
| state.set("width", 32) | ||
| _, state = vae_encoder(pipe, state) |
There was a problem hiding this comment.
same here, can you use you turn the block into pipeline to run?
3 participants
What does this PR do?
Adds support for Cosmos 3 Edge Generator