(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True, offload_model:bool=False)
| 715 | |
| 716 | |
| 717 | def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True, offload_model:bool=False): |
| 718 | input_is_list = isinstance(x, list) |
| 719 | x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent) |
| 720 | time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1) |
| 721 | |
| 722 | if input_img_latents is not None: |
| 723 | input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True) |
| 724 | if input_ids is not None: |
| 725 | condition_embeds = self.llm.embed_tokens(input_ids).clone() |
| 726 | input_img_inx = 0 |
| 727 | for b_inx in input_image_sizes.keys(): |
| 728 | for start_inx, end_inx in input_image_sizes[b_inx]: |
| 729 | condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx] |
| 730 | input_img_inx += 1 |
| 731 | if input_img_latents is not None: |
| 732 | assert input_img_inx == len(input_latents) |
| 733 | |
| 734 | input_emb = torch.cat([condition_embeds, time_token, x], dim=1) |
| 735 | else: |
| 736 | input_emb = torch.cat([time_token, x], dim=1) |
| 737 | output = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, offload_model=offload_model) |
| 738 | output, past_key_values = output.last_hidden_state, output.past_key_values |
| 739 | if input_is_list: |
| 740 | image_embedding = output[:, -max(num_tokens):] |
| 741 | time_emb = self.t_embedder(timestep, dtype=x.dtype) |
| 742 | x = self.final_layer(image_embedding, time_emb) |
| 743 | latents = [] |
| 744 | for i in range(x.size(0)): |
| 745 | latent = x[i:i+1, :num_tokens[i]] |
| 746 | latent = self.unpatchify(latent, shapes[i][0], shapes[i][1]) |
| 747 | latents.append(latent) |
| 748 | else: |
| 749 | image_embedding = output[:, -num_tokens:] |
| 750 | time_emb = self.t_embedder(timestep, dtype=x.dtype) |
| 751 | x = self.final_layer(image_embedding, time_emb) |
| 752 | latents = self.unpatchify(x, shapes[0], shapes[1]) |
| 753 | |
| 754 | if return_past_key_values: |
| 755 | return latents, past_key_values |
| 756 | return latents |
| 757 | |
| 758 | |
| 759 | @torch.no_grad() |
no test coverage detected