vllm_omni.diffusion.models.audiox.pipeline_audiox ¶

pretransform = _build_audiox_oobleck()

proj_features_128 = Linear(latent_seq_len, 128)

proj_out = (
    Linear(in_ch, cond_dim)
    if in_ch != cond_dim
    else Identity()
)

audio_channels: int = 2

audio_sample_rate: int = 44100

audio_vae_adapter = AudioVaePromptAdapter(
    cond_dim=int(model_config["conditioning"]["cond_dim"]),
    latent_seq_len=int(
        cond_configs["audio_prompt"]["latent_seq_len"]
    ),
)

clip_empty_visual_feat = Parameter(
    zeros(1, _clip_out_features, _DIM), requires_grad=False
)

clip_encoder = CLIPVisionModelWithProjection(vision_config)

clip_proj = Linear(_in_features, _clip_out_features)

clip_proj_sync = Linear(240, _clip_out_features)

clip_sync_weight = Parameter(tensor(0.0))

clip_temp_pos_embedding = Parameter(
    randn(1, _VIDEO_FPS * _DURATION_SEC, _DIM)
)

clip_temp_transformer = SA_Transformer(
    _DIM, depth=4, heads=16, dim_head=64, mlp_dim=_DIM * 4
)

device = get_local_device()

diffusion_objective = 'v'

io_channels = model_config['io_channels']

maf_block = MAF_Block(
    dim=768,
    num_experts_per_modality=int(
        gate_type_config["num_experts_per_modality"]
    ),
    num_heads=int(gate_type_config["num_heads"]),
    num_fusion_layers=int(
        gate_type_config["num_fusion_layers"]
    ),
)

model = MMDiffusionTransformer(
    **(dict(diffusion_config["config"]))
)

od_config = od_config

pretransform = _build_audiox_oobleck(
    scaling_factor=float(get("scale", 1.0))
)

support_audio_output: bool = True

text_encoder = to(float16)

tokenizer = from_pretrained(t5_name)

weights_sources = [
    ComponentSource(
        model_or_path=_model_root,
        subfolder="transformer",
        revision=getattr(od_config, "revision", None),
        prefix="",
    )
]

DIM = 768

MLP_RATIO = 4.0

bypass_gate_a = Parameter(tensor(-10.0))

bypass_gate_t = Parameter(tensor(-10.0))

bypass_gate_v = Parameter(tensor(-10.0))

cross_block = _MAFCrossAttentionBlock(dim, num_heads)

fusion_blocks = ModuleList(
    [
        (_MAFFusionBlock(dim, num_heads, mlp_ratio))
        for _ in (range(num_fusion_layers))
    ]
)

gating_network = Sequential(
    Linear(dim * 3, dim), GELU(), Linear(dim, 3), Sigmoid()
)

norm1 = LayerNorm(dim)

norm_a2 = LayerNorm(dim)

norm_t2 = LayerNorm(dim)

norm_v2 = LayerNorm(dim)

unified_experts = Parameter(randn(total_experts, dim))

heads = heads

scale = dim_head ** -0.5

to_out = (
    Sequential(Linear(inner_dim, dim), Dropout(0.0))
    if project_out
    else Identity()
)

to_qkv = Linear(dim, inner_dim * 3, bias=False)

net = Sequential(
    Linear(dim, hidden_dim),
    GELU(),
    Dropout(0.0),
    Linear(hidden_dim, dim),
    Dropout(0.0),
)

fn = fn

norm = LayerNorm(dim)

layers = ModuleList([])

norm = LayerNorm(dim)