vllm_omni.diffusion.models.cosmos3 ¶

color_format: str = 'RGB'

current_sigma

current_step_index

device = get_local_device()

do_classifier_free_guidance

dtype = od_config.dtype

guidance_scale

is_distilled_model = (
    scheduler_config.get("_class_name")
    == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS
)

is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer

num_timesteps

od_config = od_config

remap_checkpoint_key = _remap_ckpt_key

scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    scheduler_config, stochastic_sampling=True
)

support_image_input: bool = True

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

transformer = transformer_cls(
    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="transformer",
        revision=None,
        prefix="transformer.",
        fall_back_to_pt=True,
    )
]

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 = nn.Parameter(
    torch.zeros(self.hidden_size, dtype=dtype)
)

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

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

audio_modality_embed = nn.Parameter(
    torch.zeros(self.hidden_size)
)

audio_proj_in = nn.Linear(self.sound_dim, self.hidden_size)

audio_proj_out = nn.Linear(
    self.hidden_size, self.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 = nn.ModuleList(
    [
        Cosmos3GenDecoderLayer(
            layer_idx=i,
            hidden_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            num_attention_heads=self.num_attention_heads,
            num_key_value_heads=self.num_key_value_heads,
            head_dim=self.head_dim,
            rms_norm_eps=self.rms_norm_eps,
            quant_config=quant_config,
            mlp_cls=self._gen_mlp_cls,
            qk_norm=self.qk_norm_for_diffusion,
            prefix=f"gen_layers.{i}",
        )
        for i in range(self.num_hidden_layers)
    ]
)

gen_sp_gather = nn.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 = self._language_model_cls(
    hidden_size=self.hidden_size,
    intermediate_size=self.intermediate_size,
    num_hidden_layers=self.num_hidden_layers,
    num_attention_heads=self.num_attention_heads,
    num_key_value_heads=self.num_key_value_heads,
    head_dim=self.head_dim,
    vocab_size=self.vocab_size,
    rms_norm_eps=self.rms_norm_eps,
    rope_theta=self.rope_theta,
    mrope_section=self.mrope_section,
    quant_config=quant_config,
    prefix="language_model",
    **self._language_model_kwargs(),
)

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)
)

mixed_precision_runtime: (
    Cosmos3MixedPrecisionRuntime | None
) = None

mrope_section = self._resolve_mrope_section(model_config)

norm_moe_gen = RMSNorm(
    self.hidden_size, eps=self.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 = (
    self.latent_patch_size**2 * self.latent_channel_size
)

proj_in = nn.Linear(
    self.patch_latent_dim, self.hidden_size
)

proj_out = nn.Linear(
    self.hidden_size, self.patch_latent_dim
)

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

rms_norm_eps = self._resolve_rms_norm_eps(model_config)

rope_theta = self._resolve_rope_theta(model_config)

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(self.hidden_size)

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

use_und_k_norm_for_gen = _tf_config_get(
    model_config, "use_und_k_norm_for_gen", None
)

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