vllm_omni.diffusion.models.t5_encoder ¶

Tensor-parallel T5 encoder model.

Modules:

Name	Description
`t5_encoder`
`t5_gemma_encoder`

T5EncoderModel ¶

Bases: Module

T5 encoder model applying upstream vLLM layers

config = config

device: device

dtype: dtype

encoder = T5Stack(
    config, shared, prefix=f"{prefix}.encoder"
)

prefix = prefix

shared = VocabParallelEmbedding(vocab_size, d_model)

embed_input_ids(input_ids: Tensor) -> Tensor

forward(
    input_ids: Tensor, attention_mask: Tensor | None = None
) -> tuple[Tensor, ...]

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