| 1 | import torch |
| 2 | |
| 3 | |
| 4 | class PixArtAlphaTextProjection(torch.nn.Module): |
| 5 | """ |
| 6 | Projects caption embeddings using dual linear layers. |
| 7 | Flow: linear_1 → activation → linear_2 |
| 8 | Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py |
| 9 | """ |
| 10 | |
| 11 | def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): |
| 12 | super().__init__() |
| 13 | if out_features is None: |
| 14 | out_features = hidden_size |
| 15 | self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) |
| 16 | if act_fn == "gelu_tanh": |
| 17 | self.act_1 = torch.nn.GELU(approximate="tanh") |
| 18 | elif act_fn == "silu": |
| 19 | self.act_1 = torch.nn.SiLU() |
| 20 | else: |
| 21 | raise ValueError(f"Unknown activation function: {act_fn}") |
| 22 | self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) |
| 23 | |
| 24 | def forward(self, caption: torch.Tensor) -> torch.Tensor: |
| 25 | hidden_states = self.linear_1(caption) |
| 26 | hidden_states = self.act_1(hidden_states) |
| 27 | hidden_states = self.linear_2(hidden_states) |
| 28 | return hidden_states |
| 29 | |
| 30 | |
| 31 | def create_caption_projection(transformer_config: dict, audio: bool = False) -> PixArtAlphaTextProjection: |
| 32 | """Create a caption projection for the transformer (V1/19B only).""" |
| 33 | caption_channels = transformer_config["caption_channels"] |
| 34 | if audio: |
| 35 | inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"] |
| 36 | else: |
| 37 | inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"] |
| 38 | return PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) |
| 39 |