vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 ¶

Bases: Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin

Cosmos3 text/image/video/sound/action pipeline.

Architecture: Mixture-of-Transformers with Qwen3-VL backbone. - Understanding pathway: causal self-attention on text (runs once, K/V cached) - Generation pathway: cross-attention on visual latents and optional transfer-control, action, and sound latents (runs each step)

Supports T2V, I2V, V2V, T2I, transfer, sound-enabled video, and action generation from the same class. Mode is selected at runtime:

• T2I when prompt["modalities"] contains "image". Latent T-dim is forced to 1, T2I-specific scheduler defaults are applied (50 steps, flow_shift=3.0, guidance_interval=[400, 1000]), the duration template is suppressed, and post-process emits PIL images.
• I2V when the request supplies a preprocessed image via multi_modal_data['image'] (handled by :func:get_cosmos3_pre_process_func) and the requested output modality is not image. Frame 0 of the initial latent is set to the VAE-encoded conditioning image, frame-0 noise predictions are masked to zero, and the clean image latent is re-injected at frame 0 after each scheduler step.
• V2V when the request supplies a preprocessed video via multi_modal_data['video'] without an action mode. Explicit latent frame indexes are kept clean with noisy_frame_mask and re-injected after each scheduler step.
• Transfer when edge, blur, depth, seg, or wsm hints are supplied. Transfer is video-output only and cannot be combined with sound or action generation.
• Sound-enabled video when generate_sound or sound_gen is true. Sound is generated from sound latents, not from multi_modal_data['audio']; T2I, transfer, and action+sound are rejected.
• Action generation when action_mode is provided. policy and forward_dynamics require an image or video input; inverse_dynamics requires video input. Action predictions are returned in custom_output. RoboLab/OpenPI observations in extra_args['robot_obs'] or extra_args['observation'] return action-only custom output.
• T2V otherwise (default video generation).

color_format: str = 'RGB'

device = get_local_device()

do_classifier_free_guidance

dtype = od_config.dtype

guidance_scale

num_timesteps

od_config = od_config

scheduler = UniPCMultistepScheduler.from_pretrained(
    model_path,
    subfolder="scheduler",
    local_files_only=local_files_only,
)

support_image_input: bool = True

tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    subfolder="text_tokenizer",
    local_files_only=local_files_only,
)

transformer = Cosmos3VFMTransformer(
    od_config=od_config,
    temporal_compression_factor=self.vae_scale_factor_temporal,
    sound_gen=sound_gen,
    sound_dim=sound_dim,
    sound_latent_fps=sound_latent_fps,
)

vae = DistributedAutoencoderKLWan.from_pretrained(
    model_path,
    subfolder="vae",
    torch_dtype=self.dtype,
    local_files_only=local_files_only,
).to(self.device)

vae_scale_factor_spatial = getattr(
    self.vae.config, "scale_factor_spatial", 16
)

vae_scale_factor_temporal = int(
    self.vae.config.scale_factor_temporal
)

video_processor = VideoProcessor(
    vae_scale_factor=self.vae_scale_factor_spatial
)

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=model_path,
        subfolder=None,
        revision=None,
        prefix="transformer.",
        fall_back_to_pt=True,
        allow_patterns_overrides=[
            "transformer/*.safetensors"
        ],
    )
]

combine_multi_branch_cfg_noise(
    predictions: list[Tensor | tuple[Tensor, ...]],
    true_cfg_scale: float | dict[str, float],
    cfg_normalize: bool = False,
) -> Tensor | tuple[Tensor, ...]

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    cond_ids: Tensor,
    cond_mask: Tensor,
    uncond_ids: Tensor,
    uncond_mask: Tensor,
    guidance_scale: float,
    shared_kwargs: dict,
    *,
    action_latents: Tensor | None = None,
    action_velocity_mask: Tensor | None = None,
    action_condition_latents: Tensor | None = None,
    sound_latents: Tensor | None = None,
    velocity_mask: Tensor | None = None,
    image_latent: Tensor | None = None,
    condition_latents: Tensor | None = None,
    guidance_interval: tuple[float, float] | None = None,
    raw_action_dim: int | None = None,
    scheduler: Any | None = None,
) -> Tensor | tuple[Tensor, ...]

diffuse_transfer(
    latents: Tensor,
    timesteps: Tensor,
    cond_ids: Tensor,
    cond_mask: Tensor,
    uncond_ids: Tensor,
    uncond_mask: Tensor,
    guidance_scale: float,
    control_guidance: float,
    control_guidance_interval: tuple[float, float] | None,
    control_latents: list[Tensor],
    shared_kwargs: dict[str, Any],
    *,
    velocity_mask: Tensor,
    condition_latents: Tensor,
    guidance_interval: tuple[float, float] | None = None,
) -> Tensor

forward(req: DiffusionRequestBatch) -> DiffusionOutput

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

predict_noise(**kwargs) -> Tensor | tuple[Tensor, ...]

reference_video_decode_spec(
    *,
    num_frames: int | None = None,
    extra_args: dict[str, Any] | None = None,
) -> ReferenceVideoDecodeSpec