Class ConditionersModule

kleidiai-examples/audiogen/scripts/model.py:150–224

Conditioners Module. Takes the T5 encoder and seconds_total conditioner, and returns the cross-attention inputs and global conditioning inputs.

        return float_embeds, torch.ones(float_embeds.shape[0], 1)

class ConditionersModule(torch.nn.Module):
    """Conditioners Module. Takes the T5 encoder and seconds_total conditioner,
    and returns the cross-attention inputs and global conditioning inputs.
    """

    def __init__(
        self,
        sao_t5_cond: torch.nn.Module,
        sao_seconds_total_cond: torch.nn.Module,
        dtype: torch.dtype = torch.float
    ):
        super().__init__()
        self.sao_t5 = sao_t5_cond
        self.sao_seconds_total_cond = ExportableNumberConditioner(
            sao_seconds_total_cond
        )
        self.dtype = dtype

        # Use float
        self.sao_t5 = force_t5_conditioner_float32(self.sao_t5.to("cpu"))
        self.sao_t5 = self.sao_t5.to(dtype).eval().requires_grad_(False)
        self.sao_seconds_total_cond = self.sao_seconds_total_cond.to(dtype=dtype)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        seconds_total: torch.Tensor,
    ):
        # Get the projections and conditioner results
        with torch.no_grad():
            t5_embeddings = self.sao_t5.model(
                input_ids=input_ids, attention_mask=attention_mask
            )["last_hidden_state"]
            # Resize the embeddings and attention mask to 64 to match DiT model
            t5_embeddings = t5_embeddings[:, :64, :]
            attention_mask = attention_mask[:, :64]
            # Get the T5 projections
            t5_proj = self.sao_t5.proj_out(t5_embeddings).to(dtype=self.dtype)
            t5_proj = t5_proj * attention_mask.unsqueeze(-1).to(dtype=self.dtype)
            t5_mask = attention_mask

        # Get seconds_total conditioner results
        seconds_total_embedding, seconds_total_mask = self.sao_seconds_total_cond(
            seconds_total
        )

        # Concatenate all cross-attention inputs (t5_embedding, seconds_total) over the sequence dimension
        # Assumes that the cross-attention inputs are of shape (batch, seq, channels)
        cross_attention_input = torch.cat(
            [
                t5_proj,
                seconds_total_embedding,
            ],
            dim=1,
        )
        cross_attention_masks = torch.cat(
            [