vllm_omni.diffusion.models.cosmos3 ¶

color_format: str = 'RGB'

device = get_local_device()

do_classifier_free_guidance

dtype = dtype

guidance_scale

num_timesteps

od_config = od_config

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

support_image_input: bool = True

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

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

vae = to(device)

vae_scale_factor_spatial = getattr(
    config, "scale_factor_spatial", 16
)

vae_scale_factor_temporal = int(scale_factor_temporal)

video_processor = VideoProcessor(
    vae_scale_factor=vae_scale_factor_spatial
)

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

action_dim = int(
    action_dim_value if action_dim_value is not None else 64
)

action_gen = (
    _as_bool(action_gen_value)
    if action_gen_value is not None
    else False
)

action_modality_embed = Parameter(
    zeros(hidden_size, dtype=dtype)
)

action_proj_in = DomainAwareLinear(
    action_dim,
    hidden_size,
    num_embodiment_domains,
    dtype=dtype,
)

action_proj_out = DomainAwareLinear(
    hidden_size,
    action_dim,
    num_embodiment_domains,
    dtype=dtype,
)

audio_modality_embed = Parameter(zeros(hidden_size))

audio_proj_in = Linear(sound_dim, hidden_size)

audio_proj_out = Linear(hidden_size, sound_dim)

base_fps = float(
    _tf_config_get(model_config, "base_fps", 24.0)
)

cached_freqs_gen: tuple[Tensor, Tensor] | None = None

cached_kv: list[tuple[Tensor, Tensor]] | None = None

device: device

enable_fps_modulation = bool(
    _tf_config_get(
        model_config, "enable_fps_modulation", True
    )
)

gen_layers = ModuleList(
    [
        (
            Cosmos3GenDecoderLayer(
                layer_idx=i,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_attention_heads=num_attention_heads,
                num_key_value_heads=num_key_value_heads,
                head_dim=head_dim,
                rms_norm_eps=rms_norm_eps,
                quant_config=quant_config,
                prefix=f"gen_layers.{i}",
            )
        )
        for i in (range(num_hidden_layers))
    ]
)

gen_sp_gather = Identity()

gen_sp_prepare = Cosmos3GenSPPrepare()

head_dim = int(
    _tf_config_get(model_config, "head_dim", 128)
)

hidden_size = int(
    _tf_config_get(model_config, "hidden_size", 4096)
)

intermediate_size = int(
    _tf_config_get(model_config, "intermediate_size", 12288)
)

language_model = Cosmos3LanguageModel(
    hidden_size=hidden_size,
    intermediate_size=intermediate_size,
    num_hidden_layers=num_hidden_layers,
    num_attention_heads=num_attention_heads,
    num_key_value_heads=num_key_value_heads,
    head_dim=head_dim,
    vocab_size=vocab_size,
    rms_norm_eps=rms_norm_eps,
    rope_theta=rope_theta,
    mrope_section=mrope_section,
    quant_config=quant_config,
    prefix="language_model",
)

latent_channel_size = int(
    _tf_config_get(model_config, "latent_channel", 48)
)

latent_patch_size = int(
    _tf_config_get(model_config, "latent_patch_size", 2)
)

mrope_section = list(get('mrope_section', [24, 20, 20]))

norm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps)

num_attention_heads = int(
    _tf_config_get(model_config, "num_attention_heads", 32)
)

num_embodiment_domains = int(
    _od_config_get(od_config, "num_embodiment_domains", 32)
)

num_hidden_layers = int(
    _tf_config_get(model_config, "num_hidden_layers", 36)
)

num_key_value_heads = int(
    _tf_config_get(model_config, "num_key_value_heads", 8)
)

packed_modules_mapping = {}

patch_latent_dim = (
    latent_patch_size**2 * latent_channel_size
)

proj_in = Linear(patch_latent_dim, hidden_size)

proj_out = Linear(hidden_size, patch_latent_dim)

rms_norm_eps = float(
    _tf_config_get(model_config, "rms_norm_eps", 1e-06)
)

rope_theta = float(
    _tf_config_get(model_config, "rope_theta", 5000000)
)

sound_dim = sound_dim

sound_gen = sound_gen

sound_latent_fps = sound_latent_fps

temporal_compression_factor = int(
    temporal_compression_factor
)

temporal_compression_factor_sound = int(
    _tf_config_get(
        model_config, "temporal_compression_factor_sound", 1
    )
)

temporal_modality_margin = int(
    _tf_config_get(
        model_config,
        "unified_3d_mrope_temporal_modality_margin",
        15000,
    )
)

time_embedder = TimestepEmbedder(
    hidden_size, target_dtype=dtype
)

timestep_scale = float(
    _tf_config_get(model_config, "timestep_scale", 0.001)
)

vocab_size = int(
    _tf_config_get(model_config, "vocab_size", 151936)
)