| 322 | logger.info(f"Sample inputs spec: {sample_text_encoder_inputs_spec}") |
| 323 | |
| 324 | class TextEncoder(nn.Module): |
| 325 | |
| 326 | def __init__(self, with_hidden_states_for_layer=None): |
| 327 | super().__init__() |
| 328 | self.text_encoder = text_encoder |
| 329 | self.with_hidden_states_for_layer = with_hidden_states_for_layer |
| 330 | |
| 331 | def forward(self, input_ids): |
| 332 | if self.with_hidden_states_for_layer is not None: |
| 333 | output = self.text_encoder(input_ids, output_hidden_states=True) |
| 334 | hidden_embeds = output.hidden_states[self.with_hidden_states_for_layer] |
| 335 | if "text_embeds" in output: |
| 336 | return (hidden_embeds, output.text_embeds) |
| 337 | else: |
| 338 | return (hidden_embeds, output.pooler_output) |
| 339 | else: |
| 340 | return self.text_encoder(input_ids, return_dict=False) |
| 341 | |
| 342 | # SD XL uses the hidden states after the encoder layers from both encoders, |
| 343 | # and the pooled `text_embeds` output of the second encoder. |
no outgoing calls
no test coverage detected