Conditioners Module. Takes the T5 encoder and seconds_total conditioner, and returns the cross-attention inputs and global conditioning inputs.
| 123 | return float_embeds, torch.ones(float_embeds.shape[0], 1) |
| 124 | |
| 125 | class ConditionersModule(torch.nn.Module): |
| 126 | """Conditioners Module. Takes the T5 encoder and seconds_total conditioner, |
| 127 | and returns the cross-attention inputs and global conditioning inputs. |
| 128 | """ |
| 129 | |
| 130 | def __init__( |
| 131 | self, |
| 132 | sao_t5_cond: torch.nn.Module, |
| 133 | sao_seconds_total_cond: torch.nn.Module, |
| 134 | dtype: torch.dtype = torch.float |
| 135 | ): |
| 136 | super().__init__() |
| 137 | self.sao_t5 = sao_t5_cond |
| 138 | self.sao_seconds_total_cond = ExportableNumberConditioner( |
| 139 | sao_seconds_total_cond |
| 140 | ) |
| 141 | self.dtype = dtype |
| 142 | |
| 143 | # Use float |
| 144 | self.sao_t5 = ( |
| 145 | self.sao_t5.to("cpu").to(dtype).eval().requires_grad_(False) |
| 146 | ) |
| 147 | self.sao_seconds_total_cond = self.sao_seconds_total_cond.to(dtype=dtype) |
| 148 | |
| 149 | def forward( |
| 150 | self, |
| 151 | input_ids: torch.Tensor, |
| 152 | attention_mask: torch.Tensor, |
| 153 | seconds_total: torch.Tensor, |
| 154 | ): |
| 155 | # Get the projections and conditioner results |
| 156 | with torch.no_grad(): |
| 157 | t5_embeddings = self.sao_t5.model( |
| 158 | input_ids=input_ids, attention_mask=attention_mask |
| 159 | )["last_hidden_state"] |
| 160 | # Resize the embeddings and attention mask to 64 to match DiT model |
| 161 | t5_embeddings = t5_embeddings[:, :64, :] |
| 162 | attention_mask = attention_mask[:, :64] |
| 163 | # Get the T5 projections |
| 164 | t5_proj = self.sao_t5.proj_out(t5_embeddings).to(dtype=self.dtype) |
| 165 | t5_proj = t5_proj * attention_mask.unsqueeze(-1).to(dtype=self.dtype) |
| 166 | t5_mask = attention_mask |
| 167 | |
| 168 | # Get seconds_total conditioner results |
| 169 | seconds_total_embedding, seconds_total_mask = self.sao_seconds_total_cond( |
| 170 | seconds_total |
| 171 | ) |
| 172 | |
| 173 | # Concatenate all cross-attention inputs (t5_embedding, seconds_total) over the sequence dimension |
| 174 | # Assumes that the cross-attention inputs are of shape (batch, seq, channels) |
| 175 | cross_attention_input = torch.cat( |
| 176 | [ |
| 177 | t5_proj, |
| 178 | seconds_total_embedding, |
| 179 | ], |
| 180 | dim=1, |
| 181 | ) |
| 182 | cross_attention_masks = torch.cat( |
no outgoing calls
no test coverage detected