Projects text embeddings. Also handles dropout for classifier-free guidance. Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
| 68 | |
| 69 | |
| 70 | class TextProjection(nn.Module): |
| 71 | """ |
| 72 | Projects text embeddings. Also handles dropout for classifier-free guidance. |
| 73 | |
| 74 | Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py |
| 75 | """ |
| 76 | |
| 77 | def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): |
| 78 | factory_kwargs = {"dtype": dtype, "device": device} |
| 79 | super().__init__() |
| 80 | self.linear_1 = nn.Linear( |
| 81 | in_features=in_channels, |
| 82 | out_features=hidden_size, |
| 83 | bias=True, |
| 84 | **factory_kwargs, |
| 85 | ) |
| 86 | self.act_1 = act_layer() |
| 87 | self.linear_2 = nn.Linear( |
| 88 | in_features=hidden_size, |
| 89 | out_features=hidden_size, |
| 90 | bias=True, |
| 91 | **factory_kwargs, |
| 92 | ) |
| 93 | |
| 94 | def forward(self, caption): |
| 95 | hidden_states = self.linear_1(caption) |
| 96 | hidden_states = self.act_1(hidden_states) |
| 97 | hidden_states = self.linear_2(hidden_states) |
| 98 | return hidden_states |
| 99 | |
| 100 | |
| 101 | class TimestepEmbedder(nn.Module): |