vllm_omni.diffusion.models.audiox.pipeline_audiox ¶

pretransform = _build_audiox_oobleck()

proj_features_128 = nn.Linear(latent_seq_len, 128)

proj_out = (
    nn.Linear(in_ch, cond_dim)
    if in_ch != cond_dim
    else nn.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 = nn.Parameter(
    torch.zeros(1, self._clip_out_features, _DIM),
    requires_grad=False,
)

clip_encoder = CLIPVisionModelWithProjection(
    clip_config.vision_config
)

clip_proj = nn.Linear(_in_features, self._clip_out_features)

clip_proj_sync = nn.Linear(240, self._clip_out_features)

clip_sync_weight = nn.Parameter(torch.tensor(0.0))

clip_temp_pos_embedding = nn.Parameter(
    torch.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(
        model_config["pretransform"].get("scale", 1.0)
    )
)

support_audio_output: bool = True

text_encoder = (
    T5EncoderModel(t5_config)
    .train(False)
    .requires_grad_(False)
    .to(torch.float16)
)

tokenizer = T5TokenizerFast.from_pretrained(t5_name)

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

DIM = 768

MLP_RATIO = 4.0

bypass_gate_a = nn.Parameter(torch.tensor(-10.0))

bypass_gate_t = nn.Parameter(torch.tensor(-10.0))

bypass_gate_v = nn.Parameter(torch.tensor(-10.0))

cross_block = _MAFCrossAttentionBlock(dim, num_heads)

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

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

norm1 = nn.LayerNorm(dim)

norm_a2 = nn.LayerNorm(dim)

norm_t2 = nn.LayerNorm(dim)

norm_v2 = nn.LayerNorm(dim)

unified_experts = nn.Parameter(
    torch.randn(total_experts, dim)
)

heads = heads

scale = dim_head ** -0.5

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

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

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

fn = fn

norm = nn.LayerNorm(dim)

layers = nn.ModuleList([])

norm = nn.LayerNorm(dim)