mstar.model.cosmos3.tests.pipeline

mstar.model.cosmos3.tests.pipeline#

Fused generation pipeline for Cosmos3-Nano (text/image-to-image/video).

Runs the generator in one fused forward per denoising step (text + vision together), using mstar’s DiT forward + packing and the imported diffusers UniPC scheduler + Wan VAE. Intentionally simple (batch 1, sequential CFG); not the served path. Produces the same image/video as the diffusers Cosmos3OmniPipeline on a fixed seed/prompt.

num_frames == 1 is text-to-image; num_frames > 1 is text-to-video, and passing image anchors frame 0 to a conditioning frame (image-to-video).

Classes

Cosmos3Pipeline(transformer, vae, scheduler, ...)

Fused t2i / t2v / i2v pipeline for Cosmos3-Nano.
class mstar.model.cosmos3.tests.pipeline.Cosmos3Pipeline(transformer, vae, scheduler, tokenizer, config, device, dtype=torch.bfloat16)[source]#

Bases: object

Fused t2i / t2v / i2v pipeline for Cosmos3-Nano.

classmethod from_model(model, device, dtype=torch.bfloat16)[source]#

Build from a loaded Cosmos3Model (DiT + Wan VAE) + imported UniPC.

generate_action(*, prompt, mode, domain_id, action_chunk_size, raw_action_dim, video=None, video_latents=None, action=None, num_frames=None, height=256, width=256, fps=24.0, action_fps=None, num_inference_steps=30, guidance_scale=1.0, flow_shift=None, negative_prompt='', generator=None, cond_ids=None, uncond_ids=None, return_video=False)[source]#

Joint video+action generation (forward_dynamics / inverse_dynamics / policy).

The conditioning video is VAE-encoded to clean anchor frames per mode (all frames for inverse-dynamics; frame 0 for forward-dynamics / policy). Action tokens are clean conditioning for forward-dynamics, else noisy and predicted. Returns the predicted action [1, action_chunk_size, raw_action_dim] (and the decoded video when return_video).

Parameters: