clip

class CLIPVisionEmbeddings(nn.Module):
def __init__(self, config: CLIPVisionConfig):
    super().__init__()
    self.config = config
    self.embed_dim = config.hidden_size
    self.image_size = config.image_size
    self.patch_size = config.patch_size
    assert self.image_size % self.patch_size == 0

    self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))

    self.patch_embedding = nn.Conv2d(
        in_channels=config.num_channels,
        out_channels=self.embed_dim,
        kernel_size=self.patch_size,
        stride=self.patch_size,
        bias=False,
    )

    self.num_patches = (self.image_size // self.patch_size)**2
    self.num_positions = self.num_patches + 1
    self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
    self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)

def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
    batch_size = pixel_values.shape[0]
    target_dtype = self.patch_embedding.weight.dtype
    patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]
    patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

    class_embeds = self.class_embedding.expand(batch_size, 1, -1)
    embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
    embeddings = embeddings + self.position_embedding(self.position_ids)

    return embeddings</code></pre></details>
fastvideo.models.encoders.clip.CLIPVisionEmbeddings.forward
forward(pixel_values: torch.Tensor) -> torch.Tensor
Source
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
    batch_size = pixel_values.shape[0]
    target_dtype = self.patch_embedding.weight.dtype
    patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]
    patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
class_embeds = self.class_embedding.expand(batch_size, 1, -1)
embeddings = torch.cat(&#91;class_embeds, patch_embeds&#93;, dim=1)
embeddings = embeddings + self.position_embedding(self.position_ids)

return embeddings</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextEmbeddings
CLIPTextEmbeddings(config: CLIPTextConfig)
Source
class CLIPTextEmbeddings(nn.Module):
def __init__(self, config: CLIPTextConfig):
    super().__init__()
    self.config = config
    embed_dim = config.hidden_size

    self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
    self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)

    # position_ids (1, len position emb) is contiguous in memory and exported when serialized
    self.register_buffer(&quot;position_ids&quot;,
                         torch.arange(config.max_position_embeddings).expand((1, -1)),
                         persistent=False)

def forward(
    self,
    input_ids: torch.LongTensor | None = None,
    position_ids: torch.LongTensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
) -&gt; torch.Tensor:
    if input_ids is not None:
        seq_length = input_ids.shape&#91;-1&#93;
    elif inputs_embeds is not None:
        seq_length = inputs_embeds.shape&#91;-2&#93;
    else:
        raise ValueError(&quot;Either input_ids or inputs_embeds must be provided.&quot;)

    max_position_embedding = self.position_embedding.weight.shape&#91;0&#93;

    if seq_length &gt; max_position_embedding:
        raise ValueError(f&quot;Sequence length must be less than max_position_embeddings (got `sequence length`: &quot;
                         f&quot;&#123;seq_length&#125; and max_position_embeddings: &#123;max_position_embedding&#125;&quot;)

    if position_ids is None:
        position_ids = self.position_ids&#91;:, :seq_length&#93;

    if inputs_embeds is None:
        inputs_embeds = self.token_embedding(input_ids)

    position_embeddings = self.position_embedding(position_ids)
    embeddings = inputs_embeds + position_embeddings

    return embeddings</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextEmbeddings.forward
forward(input_ids: torch.LongTensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    input_ids: torch.LongTensor | None = None,
    position_ids: torch.LongTensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
) -> torch.Tensor:
    if input_ids is not None:
        seq_length = input_ids.shape[-1]
    elif inputs_embeds is not None:
        seq_length = inputs_embeds.shape[-2]
    else:
        raise ValueError("Either input_ids or inputs_embeds must be provided.")
max_position_embedding = self.position_embedding.weight.shape&#91;0&#93;

if seq_length &gt; max_position_embedding:
    raise ValueError(f&quot;Sequence length must be less than max_position_embeddings (got `sequence length`: &quot;
                     f&quot;&#123;seq_length&#125; and max_position_embeddings: &#123;max_position_embedding&#125;&quot;)

if position_ids is None:
    position_ids = self.position_ids&#91;:, :seq_length&#93;

if inputs_embeds is None:
    inputs_embeds = self.token_embedding(input_ids)

position_embeddings = self.position_embedding(position_ids)
embeddings = inputs_embeds + position_embeddings

return embeddings</code></pre></details>
fastvideo.models.encoders.clip.CLIPAttention
CLIPAttention(config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')
Multi-headed attention from 'Attention Is All You Need' paper
Source
class CLIPAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""
def __init__(
    self,
    config: CLIPVisionConfig | CLIPTextConfig,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
):
    super().__init__()
    self.config = config
    self.embed_dim = config.hidden_size
    self.num_heads = config.num_attention_heads
    self.head_dim = self.embed_dim // self.num_heads
    if self.head_dim * self.num_heads != self.embed_dim:
        raise ValueError(&quot;embed_dim must be divisible by num_heads &quot;
                         f&quot;(got `embed_dim`: &#123;self.embed_dim&#125; and `num_heads`:&quot;
                         f&quot; &#123;self.num_heads&#125;).&quot;)
    self.scale = self.head_dim**-0.5 if config.enable_scale else None
    self.dropout = config.attention_dropout

    self.qkv_proj = QKVParallelLinear(
        hidden_size=self.embed_dim,
        head_size=self.head_dim,
        total_num_heads=self.num_heads,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.qkv_proj&quot;,
    )

    self.out_proj = RowParallelLinear(
        input_size=self.embed_dim,
        output_size=self.embed_dim,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.out_proj&quot;,
    )

    self.tp_size = get_tp_world_size()
    self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

    self.attn = LocalAttention(self.num_heads_per_partition,
                               self.head_dim,
                               self.num_heads_per_partition,
                               softmax_scale=self.scale,
                               causal=config.is_causal,
                               supported_attention_backends=config._supported_attention_backends)

def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
    return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

def forward(
    self,
    hidden_states: torch.Tensor,
):
    &quot;&quot;&quot;Input shape: Batch x Time x Channel&quot;&quot;&quot;

    qkv_states, _ = self.qkv_proj(hidden_states)
    query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
    # use flash_attn_func
    query_states = query_states.reshape(query_states.shape&#91;0&#93;, query_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                        self.head_dim)
    key_states = key_states.reshape(key_states.shape&#91;0&#93;, key_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                    self.head_dim)
    value_states = value_states.reshape(value_states.shape&#91;0&#93;, value_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                        self.head_dim)
    attn_output = self.attn(query_states, key_states, value_states)
    attn_output = attn_output.reshape(attn_output.shape&#91;0&#93;, attn_output.shape&#91;1&#93;,
                                      self.num_heads_per_partition * self.head_dim)
    attn_output, _ = self.out_proj(attn_output)

    return attn_output, None</code></pre></details>
fastvideo.models.encoders.clip.CLIPAttention.forward
forward(hidden_states: torch.Tensor)
Input shape: Batch x Time x Channel
Source
def forward(
    self,
    hidden_states: torch.Tensor,
):
    """Input shape: Batch x Time x Channel"""
qkv_states, _ = self.qkv_proj(hidden_states)
query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
# use flash_attn_func
query_states = query_states.reshape(query_states.shape&#91;0&#93;, query_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                    self.head_dim)
key_states = key_states.reshape(key_states.shape&#91;0&#93;, key_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                self.head_dim)
value_states = value_states.reshape(value_states.shape&#91;0&#93;, value_states.shape&#91;1&#93;, self.num_heads_per_partition,
                                    self.head_dim)
attn_output = self.attn(query_states, key_states, value_states)
attn_output = attn_output.reshape(attn_output.shape&#91;0&#93;, attn_output.shape&#91;1&#93;,
                                  self.num_heads_per_partition * self.head_dim)
attn_output, _ = self.out_proj(attn_output)

return attn_output, None</code></pre></details>
fastvideo.models.encoders.clip.CLIPMLP
CLIPMLP(config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')
Source
class CLIPMLP(nn.Module):
def __init__(
    self,
    config: CLIPVisionConfig | CLIPTextConfig,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()
    self.config = config
    self.activation_fn = get_act_fn(config.hidden_act)
    self.fc1 = ColumnParallelLinear(config.hidden_size,
                                    config.intermediate_size,
                                    bias=True,
                                    quant_config=quant_config,
                                    prefix=f&quot;&#123;prefix&#125;.fc1&quot;)
    self.fc2 = RowParallelLinear(config.intermediate_size,
                                 config.hidden_size,
                                 bias=True,
                                 quant_config=quant_config,
                                 prefix=f&quot;&#123;prefix&#125;.fc2&quot;)

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    hidden_states, _ = self.fc1(hidden_states)
    hidden_states = self.activation_fn(hidden_states)
    hidden_states, _ = self.fc2(hidden_states)

    return hidden_states</code></pre></details>
fastvideo.models.encoders.clip.CLIPMLP.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states, _ = self.fc1(hidden_states)
    hidden_states = self.activation_fn(hidden_states)
    hidden_states, _ = self.fc2(hidden_states)
return hidden_states</code></pre></details>
fastvideo.models.encoders.clip.CLIPEncoderLayer
CLIPEncoderLayer(config: CLIPTextConfig | CLIPVisionConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')
Source
class CLIPEncoderLayer(nn.Module):
def __init__(
    self,
    config: CLIPTextConfig | CLIPVisionConfig,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()
    self.self_attn = CLIPAttention(
        config,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.self_attn&quot;,
    )
    self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
    self.mlp = CLIPMLP(config, quant_config=quant_config, prefix=f&quot;&#123;prefix&#125;.mlp&quot;)
    self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:

    residual = hidden_states

    hidden_states = self.layer_norm1(hidden_states)
    hidden_states, _ = self.self_attn(hidden_states=hidden_states)
    hidden_states = residual + hidden_states

    residual = hidden_states
    hidden_states = self.layer_norm2(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states

    return hidden_states</code></pre></details>
fastvideo.models.encoders.clip.CLIPEncoderLayer.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
residual = hidden_states

hidden_states = self.layer_norm1(hidden_states)
hidden_states, _ = self.self_attn(hidden_states=hidden_states)
hidden_states = residual + hidden_states

residual = hidden_states
hidden_states = self.layer_norm2(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states

return hidden_states</code></pre></details>
fastvideo.models.encoders.clip.CLIPEncoder
CLIPEncoder(config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, prefix: str = '')
Transformer encoder consisting of config.num_hidden_layers self
attention layers. Each layer is a [CLIPEncoderLayer].
Parameters

config: CLIPVisionConfig | CLIPTextConfig
CLIPConfig

Source
class CLIPEncoder(nn.Module):
    """
    Transformer encoder consisting of `config.num_hidden_layers` self
    attention layers. Each layer is a [`CLIPEncoderLayer`].
Args:
    config: CLIPConfig
&quot;&quot;&quot;

def __init__(
    self,
    config: CLIPVisionConfig | CLIPTextConfig,
    quant_config: QuantizationConfig | None = None,
    num_hidden_layers_override: int | None = None,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()

    self.config = config

    if num_hidden_layers_override is None:
        num_hidden_layers = config.num_hidden_layers
    else:
        num_hidden_layers = num_hidden_layers_override
    self.layers = nn.ModuleList(&#91;
        CLIPEncoderLayer(config=config, quant_config=quant_config, prefix=f&quot;&#123;prefix&#125;.layers.&#123;layer_idx&#125;&quot;)
        for layer_idx in range(num_hidden_layers)
    &#93;)

def forward(self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool) -&gt; torch.Tensor | list&#91;torch.Tensor&#93;:
    hidden_states_pool = &#91;inputs_embeds&#93;
    hidden_states = inputs_embeds

    for encoder_layer in self.layers:
        hidden_states = encoder_layer(hidden_states)
        if return_all_hidden_states:
            hidden_states_pool.append(hidden_states)
    # If we have multiple feature sample layers, we return all hidden
    # states in order and grab the ones we need by index.
    if return_all_hidden_states:
        return hidden_states_pool
    return &#91;hidden_states&#93;</code></pre></details>
fastvideo.models.encoders.clip.CLIPEncoder.forward
forward(inputs_embeds: torch.Tensor, return_all_hidden_states: bool) -> torch.Tensor | list[torch.Tensor]
Source
def forward(self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool) -> torch.Tensor | list[torch.Tensor]:
    hidden_states_pool = [inputs_embeds]
    hidden_states = inputs_embeds
for encoder_layer in self.layers:
    hidden_states = encoder_layer(hidden_states)
    if return_all_hidden_states:
        hidden_states_pool.append(hidden_states)
# If we have multiple feature sample layers, we return all hidden
# states in order and grab the ones we need by index.
if return_all_hidden_states:
    return hidden_states_pool
return &#91;hidden_states&#93;</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextTransformer
CLIPTextTransformer(config: CLIPTextConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, prefix: str = '')
Source
class CLIPTextTransformer(nn.Module):
def __init__(self,
             config: CLIPTextConfig,
             quant_config: QuantizationConfig | None = None,
             num_hidden_layers_override: int | None = None,
             prefix: str = &quot;&quot;):
    super().__init__()
    self.config = config
    embed_dim = config.hidden_size

    self.embeddings = CLIPTextEmbeddings(config)

    self.encoder = CLIPEncoder(config,
                               quant_config=quant_config,
                               num_hidden_layers_override=num_hidden_layers_override,
                               prefix=prefix)

    self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    # For `pooled_output` computation
    self.eos_token_id = config.eos_token_id

def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
) -&gt; BaseEncoderOutput:
    r&quot;&quot;&quot;
    Returns:

    &quot;&quot;&quot;
    output_hidden_states = (output_hidden_states
                            if output_hidden_states is not None else self.config.output_hidden_states)

    if input_ids is None:
        raise ValueError(&quot;You have to specify input_ids&quot;)

    input_shape = input_ids.size()
    input_ids = input_ids.view(-1, input_shape&#91;-1&#93;)

    hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

    # CLIP&#x27;s text model uses causal mask, prepare it here.
    # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
    # causal_attention_mask = _create_4d_causal_attention_mask(
    #     input_shape, hidden_states.dtype, device=hidden_states.device
    # )

    # # expand attention_mask
    # if attention_mask is not None and not self._use_flash_attention_2:
    #     raise NotImplementedError(&quot;attention_mask is not supported for CLIPTextTransformer&quot;)
    #     # &#91;bsz, seq_len&#93; -&gt; &#91;bsz, 1, tgt_seq_len, src_seq_len&#93;
    #     attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)

    encoder_outputs = self.encoder(
        inputs_embeds=hidden_states,
        # attention_mask=attention_mask,
        # causal_attention_mask=causal_attention_mask,
        # output_attentions=output_attentions,
        return_all_hidden_states=output_hidden_states,
        # return_dict=return_dict,
    )

    last_hidden_state = encoder_outputs&#91;-1&#93;
    last_hidden_state = self.final_layer_norm(last_hidden_state)

    if self.eos_token_id == 2:
        # The `eos_token_id` was incorrect before PR #24773: Let&#x27;s keep what have been done here.
        # A CLIP model with such `eos_token_id` in the config can&#x27;t work correctly with extra new tokens added
        # ------------------------------------------------------------
        # text_embeds.shape = &#91;batch_size, sequence_length, transformer.width&#93;
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        # casting to torch.int for onnx compatibility: argmax doesn&#x27;t support int64 inputs with opset 14
        pooled_output = last_hidden_state&#91;
            torch.arange(last_hidden_state.shape&#91;0&#93;, device=last_hidden_state.device),
            input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
        &#93;
    else:
        # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
        pooled_output = last_hidden_state&#91;
            torch.arange(last_hidden_state.shape&#91;0&#93;, device=last_hidden_state.device),
            # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
            # Note: we assume each sequence (along batch dim.) contains an  `eos_token_id` (e.g. prepared by the tokenizer)
            (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id).int().argmax(
                dim=-1),
        &#93;

    return BaseEncoderOutput(
        last_hidden_state=last_hidden_state,
        pooler_output=pooled_output,
        hidden_states=encoder_outputs,
        # attentions=encoder_outputs.attentions,
    )</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextTransformer.forward
forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None) -> BaseEncoderOutput
Returns:
Source
def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
) -> BaseEncoderOutput:
    r"""
    Returns:
&quot;&quot;&quot;
output_hidden_states = (output_hidden_states
                        if output_hidden_states is not None else self.config.output_hidden_states)

if input_ids is None:
    raise ValueError(&quot;You have to specify input_ids&quot;)

input_shape = input_ids.size()
input_ids = input_ids.view(-1, input_shape&#91;-1&#93;)

hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

# CLIP&#x27;s text model uses causal mask, prepare it here.
# https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
# causal_attention_mask = _create_4d_causal_attention_mask(
#     input_shape, hidden_states.dtype, device=hidden_states.device
# )

# # expand attention_mask
# if attention_mask is not None and not self._use_flash_attention_2:
#     raise NotImplementedError(&quot;attention_mask is not supported for CLIPTextTransformer&quot;)
#     # &#91;bsz, seq_len&#93; -&gt; &#91;bsz, 1, tgt_seq_len, src_seq_len&#93;
#     attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)

encoder_outputs = self.encoder(
    inputs_embeds=hidden_states,
    # attention_mask=attention_mask,
    # causal_attention_mask=causal_attention_mask,
    # output_attentions=output_attentions,
    return_all_hidden_states=output_hidden_states,
    # return_dict=return_dict,
)

last_hidden_state = encoder_outputs&#91;-1&#93;
last_hidden_state = self.final_layer_norm(last_hidden_state)

if self.eos_token_id == 2:
    # The `eos_token_id` was incorrect before PR #24773: Let&#x27;s keep what have been done here.
    # A CLIP model with such `eos_token_id` in the config can&#x27;t work correctly with extra new tokens added
    # ------------------------------------------------------------
    # text_embeds.shape = &#91;batch_size, sequence_length, transformer.width&#93;
    # take features from the eot embedding (eot_token is the highest number in each sequence)
    # casting to torch.int for onnx compatibility: argmax doesn&#x27;t support int64 inputs with opset 14
    pooled_output = last_hidden_state&#91;
        torch.arange(last_hidden_state.shape&#91;0&#93;, device=last_hidden_state.device),
        input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
    &#93;
else:
    # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
    pooled_output = last_hidden_state&#91;
        torch.arange(last_hidden_state.shape&#91;0&#93;, device=last_hidden_state.device),
        # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
        # Note: we assume each sequence (along batch dim.) contains an  `eos_token_id` (e.g. prepared by the tokenizer)
        (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id).int().argmax(
            dim=-1),
    &#93;

return BaseEncoderOutput(
    last_hidden_state=last_hidden_state,
    pooler_output=pooled_output,
    hidden_states=encoder_outputs,
    # attentions=encoder_outputs.attentions,
)</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextModel
CLIPTextModel(config: CLIPTextConfig)
Source
class CLIPTextModel(TextEncoder):
def __init__(
    self,
    config: CLIPTextConfig,
) -&gt; None:
    super().__init__(config)
    self.text_model = CLIPTextTransformer(config=config, quant_config=config.quant_config, prefix=config.prefix)

def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -&gt; BaseEncoderOutput:

    outputs: BaseEncoderOutput = self.text_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        output_hidden_states=output_hidden_states,
    )
    return outputs

def load_weights(self, weights: Iterable&#91;tuple&#91;str, torch.Tensor&#93;&#93;) -&gt; set&#91;str&#93;:

    # Define mapping for stacked parameters
    stacked_params_mapping = &#91;
        # (param_name, shard_name, shard_id)
        (&quot;qkv_proj&quot;, &quot;q_proj&quot;, &quot;q&quot;),
        (&quot;qkv_proj&quot;, &quot;k_proj&quot;, &quot;k&quot;),
        (&quot;qkv_proj&quot;, &quot;v_proj&quot;, &quot;v&quot;),
    &#93;
    params_dict = dict(self.named_parameters())
    loaded_params: set&#91;str&#93; = set()

    weights = list(weights)
    # Some checkpoints (Kandinsky6&#x27;s, for one) save CLIP&#x27;s
    # inner CLIPTextTransformer&#x27;s own state_dict() directly, without the
    # `text_model.` prefix this class&#x27;s `self.text_model` wrapper puts on
    # every parameter -- unlike other checkpoints already loaded through
    # this class (e.g. Kandinsky5&#x27;s), whose keys already include it.
    # Detect and normalize once, up front, rather than guessing per key.
    needs_text_model_prefix = bool(weights) and not any(name.startswith(&quot;text_model.&quot;) for name, _ in weights)

    for name, loaded_weight in weights:
        if needs_text_model_prefix:
            name = f&quot;text_model.&#123;name&#125;&quot;
        # Handle q_proj, k_proj, v_proj -&gt; qkv_proj mapping
        for param_name, weight_name, shard_id in stacked_params_mapping:
            if weight_name in name:
                # Replace the weight name with the parameter name
                model_param_name = name.replace(weight_name, param_name)

                if model_param_name in params_dict:
                    param = params_dict&#91;model_param_name&#93;
                    weight_loader = param.weight_loader
                    weight_loader(param, loaded_weight, shard_id)
                    loaded_params.add(model_param_name)
                break
        else:
            # Use default weight loader for all other parameters
            if name in params_dict:
                param = params_dict&#91;name&#93;
                weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
                weight_loader(param, loaded_weight)
                loaded_params.add(name)

    return loaded_params</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextModel.forward
forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs) -> BaseEncoderOutput
Source
def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -> BaseEncoderOutput:
outputs: BaseEncoderOutput = self.text_model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    position_ids=position_ids,
    output_hidden_states=output_hidden_states,
)
return outputs</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextModel.load_weights
load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]
Source
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Define mapping for stacked parameters
stacked_params_mapping = &#91;
    # (param_name, shard_name, shard_id)
    (&quot;qkv_proj&quot;, &quot;q_proj&quot;, &quot;q&quot;),
    (&quot;qkv_proj&quot;, &quot;k_proj&quot;, &quot;k&quot;),
    (&quot;qkv_proj&quot;, &quot;v_proj&quot;, &quot;v&quot;),
&#93;
params_dict = dict(self.named_parameters())
loaded_params: set&#91;str&#93; = set()

weights = list(weights)
# Some checkpoints (Kandinsky6&#x27;s, for one) save CLIP&#x27;s
# inner CLIPTextTransformer&#x27;s own state_dict() directly, without the
# `text_model.` prefix this class&#x27;s `self.text_model` wrapper puts on
# every parameter -- unlike other checkpoints already loaded through
# this class (e.g. Kandinsky5&#x27;s), whose keys already include it.
# Detect and normalize once, up front, rather than guessing per key.
needs_text_model_prefix = bool(weights) and not any(name.startswith(&quot;text_model.&quot;) for name, _ in weights)

for name, loaded_weight in weights:
    if needs_text_model_prefix:
        name = f&quot;text_model.&#123;name&#125;&quot;
    # Handle q_proj, k_proj, v_proj -&gt; qkv_proj mapping
    for param_name, weight_name, shard_id in stacked_params_mapping:
        if weight_name in name:
            # Replace the weight name with the parameter name
            model_param_name = name.replace(weight_name, param_name)

            if model_param_name in params_dict:
                param = params_dict&#91;model_param_name&#93;
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                loaded_params.add(model_param_name)
            break
    else:
        # Use default weight loader for all other parameters
        if name in params_dict:
            param = params_dict&#91;name&#93;
            weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

return loaded_params</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextModelWithProjection
CLIPTextModelWithProjection(config: CLIPTextConfig)
Source
class CLIPTextModelWithProjection(CLIPTextModel):
def __init__(self, config: CLIPTextConfig) -&gt; None:
    super().__init__(config)
    self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)

def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -&gt; BaseEncoderOutput:
    outputs = super().forward(
        input_ids=input_ids,
        position_ids=position_ids,
        attention_mask=attention_mask,
        inputs_embeds=inputs_embeds,
        output_hidden_states=output_hidden_states,
        **kwargs,
    )
    pooled = outputs.pooler_output
    if pooled is not None:
        pooled = self.text_projection(pooled)
    return BaseEncoderOutput(
        last_hidden_state=outputs.last_hidden_state,
        pooler_output=pooled,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
        attention_mask=outputs.attention_mask,
    )</code></pre></details>
fastvideo.models.encoders.clip.CLIPTextModelWithProjection.forward
forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs) -> BaseEncoderOutput
Source
def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -> BaseEncoderOutput:
    outputs = super().forward(
        input_ids=input_ids,
        position_ids=position_ids,
        attention_mask=attention_mask,
        inputs_embeds=inputs_embeds,
        output_hidden_states=output_hidden_states,
        **kwargs,
    )
    pooled = outputs.pooler_output
    if pooled is not None:
        pooled = self.text_projection(pooled)
    return BaseEncoderOutput(
        last_hidden_state=outputs.last_hidden_state,
        pooler_output=pooled,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
        attention_mask=outputs.attention_mask,
    )
fastvideo.models.encoders.clip.CLIPVisionTransformer
CLIPVisionTransformer(config: CLIPVisionConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, require_post_norm: bool | None = None, prefix: str = '')
Source
class CLIPVisionTransformer(nn.Module):
def __init__(
    self,
    config: CLIPVisionConfig,
    quant_config: QuantizationConfig | None = None,
    num_hidden_layers_override: int | None = None,
    require_post_norm: bool | None = None,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()

    self.config = config
    embed_dim = config.hidden_size

    self.embeddings = CLIPVisionEmbeddings(config)

    # NOTE: This typo of &quot;layrnorm&quot; is not fixed on purpose to match
    # the original transformers code and name of the model weights.
    self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    self.encoder = CLIPEncoder(
        config=config,
        quant_config=quant_config,
        num_hidden_layers_override=num_hidden_layers_override,
        prefix=f&quot;&#123;prefix&#125;.encoder&quot;,
    )

    num_hidden_layers = config.num_hidden_layers
    if len(self.encoder.layers) &gt; config.num_hidden_layers:
        raise ValueError(f&quot;The original encoder only has &#123;num_hidden_layers&#125; &quot;
                         f&quot;layers, but you requested &#123;len(self.encoder.layers)&#125; layers.&quot;)

    # If possible, skip post_layernorm to conserve memory
    if require_post_norm is None:
        require_post_norm = len(self.encoder.layers) == num_hidden_layers

    if require_post_norm:
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
    else:
        self.post_layernorm = None

def forward(
    self,
    pixel_values: torch.Tensor,
    feature_sample_layers: list&#91;int&#93; | None = None,
) -&gt; torch.Tensor:

    hidden_states = self.embeddings(pixel_values)
    hidden_states = self.pre_layrnorm(hidden_states)

    return_all_hidden_states = feature_sample_layers is not None

    # Produces either the last layer output or all of the hidden states,
    # depending on if we have feature_sample_layers or not
    encoder_outputs = self.encoder(inputs_embeds=hidden_states, return_all_hidden_states=return_all_hidden_states)

    if not return_all_hidden_states:
        encoder_outputs = encoder_outputs&#91;0&#93;

    # Handle post-norm (if applicable) and stacks feature layers if needed
    encoder_outputs = resolve_visual_encoder_outputs(encoder_outputs, feature_sample_layers, self.post_layernorm,
                                                     self.config.num_hidden_layers)

    return encoder_outputs</code></pre></details>
fastvideo.models.encoders.clip.CLIPVisionTransformer.forward
forward(pixel_values: torch.Tensor, feature_sample_layers: list[int] | None = None) -> torch.Tensor
Source
def forward(
    self,
    pixel_values: torch.Tensor,
    feature_sample_layers: list[int] | None = None,
) -> torch.Tensor:
hidden_states = self.embeddings(pixel_values)
hidden_states = self.pre_layrnorm(hidden_states)

return_all_hidden_states = feature_sample_layers is not None

# Produces either the last layer output or all of the hidden states,
# depending on if we have feature_sample_layers or not
encoder_outputs = self.encoder(inputs_embeds=hidden_states, return_all_hidden_states=return_all_hidden_states)

if not return_all_hidden_states:
    encoder_outputs = encoder_outputs&#91;0&#93;

# Handle post-norm (if applicable) and stacks feature layers if needed
encoder_outputs = resolve_visual_encoder_outputs(encoder_outputs, feature_sample_layers, self.post_layernorm,
                                                 self.config.num_hidden_layers)

return encoder_outputs</code></pre></details>
fastvideo.models.encoders.clip.CLIPVisionModel
CLIPVisionModel(config: CLIPVisionConfig)
Source
class CLIPVisionModel(ImageEncoder):
    config_class = CLIPVisionConfig
    main_input_name = "pixel_values"
    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
def __init__(self, config: CLIPVisionConfig) -&gt; None:
    super().__init__(config)
    self.vision_model = CLIPVisionTransformer(config=config,
                                              quant_config=config.quant_config,
                                              num_hidden_layers_override=config.num_hidden_layers_override,
                                              require_post_norm=config.require_post_norm,
                                              prefix=f&quot;&#123;config.prefix&#125;.vision_model&quot;)

def forward(
    self,
    pixel_values: torch.Tensor,
    feature_sample_layers: list&#91;int&#93; | None = None,
    **kwargs,
) -&gt; BaseEncoderOutput:
    last_hidden_state = self.vision_model(pixel_values, feature_sample_layers)
    return BaseEncoderOutput(last_hidden_state=last_hidden_state)

@property
def device(self):
    return next(self.parameters()).device

# (TODO) Add prefix argument for filtering out weights to be loaded
#        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
def load_weights(self, weights: Iterable&#91;tuple&#91;str, torch.Tensor&#93;&#93;) -&gt; set&#91;str&#93;:

    params_dict = dict(self.named_parameters())
    loaded_params: set&#91;str&#93; = set()
    layer_count = len(self.vision_model.encoder.layers)

    for name, loaded_weight in weights:
        if name.startswith(&quot;visual_projection&quot;):
            continue
        # post_layernorm is not needed in CLIPVisionModel
        if (name.startswith(&quot;vision_model.post_layernorm&quot;) and self.vision_model.post_layernorm is None):
            continue

        # omit layers when num_hidden_layers_override is set
        if name.startswith(&quot;vision_model.encoder.layers&quot;):
            layer_idx = int(name.split(&quot;.&quot;)&#91;3&#93;)
            if layer_idx &gt;= layer_count:
                continue

        for (param_name, weight_name, shard_id) in self.config.arch_config.stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)

            param = params_dict&#91;name&#93;
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            param = params_dict&#91;name&#93;
            weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
    return loaded_params</code></pre></details>
fastvideo.models.encoders.clip.CLIPVisionModel.forward
forward(pixel_values: torch.Tensor, feature_sample_layers: list[int] | None = None, **kwargs) -> BaseEncoderOutput
Source
def forward(
    self,
    pixel_values: torch.Tensor,
    feature_sample_layers: list[int] | None = None,
    **kwargs,
) -> BaseEncoderOutput:
    last_hidden_state = self.vision_model(pixel_values, feature_sample_layers)
    return BaseEncoderOutput(last_hidden_state=last_hidden_state)
fastvideo.models.encoders.clip.CLIPVisionModel.load_weights
load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]
Source
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set&#91;str&#93; = set()
layer_count = len(self.vision_model.encoder.layers)

for name, loaded_weight in weights:
    if name.startswith(&quot;visual_projection&quot;):
        continue
    # post_layernorm is not needed in CLIPVisionModel
    if (name.startswith(&quot;vision_model.post_layernorm&quot;) and self.vision_model.post_layernorm is None):
        continue

    # omit layers when num_hidden_layers_override is set
    if name.startswith(&quot;vision_model.encoder.layers&quot;):
        layer_idx = int(name.split(&quot;.&quot;)&#91;3&#93;)
        if layer_idx &gt;= layer_count:
            continue

    for (param_name, weight_name, shard_id) in self.config.arch_config.stacked_params_mapping:
        if weight_name not in name:
            continue
        name = name.replace(weight_name, param_name)

        param = params_dict&#91;name&#93;
        weight_loader = param.weight_loader
        weight_loader(param, loaded_weight, shard_id)
        break
    else:
        param = params_dict&#91;name&#93;
        weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
        weight_loader(param, loaded_weight)
    loaded_params.add(name)
return loaded_params</code></pre></details>