mstar.model.cosmos3.tests.pipeline

mstar.model.cosmos3.tests.pipeline#

Fused generation pipeline for Cosmos3 (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/Super/Edge).
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/Super/Edge).

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:
windowed_kv(cond_ids, uncond_ids, total_units, window_units, context_units, height, width, num_inference_steps, guidance_scale=6.0, fps=24.0, flow_shift=None, generator=None, page_size=128)[source]#

Block-causal windowed generation, the served kv mode’s oracle.

Per window: denoise over [text | committed context | window] with a fresh scheduler, commit the finished window’s clean K/V (no timestep embedding), then release committed frames older than context_units behind the frontier under the paged pool’s retention contract — whole pages from the first page fully past the text prefix, floor semantics with the shortfall carried to the next commit (context_units 0 retains everything). Noise draws mirror the serving path: window 0 first, then one draw per boundary from the same generator. Returns the per-window clean latents (windows carry no overlap in kv mode).

Parameters:
Return type:

list[Tensor]