(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)
| 487 | |
| 488 | |
| 489 | 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): |
| 490 | """ |
| 491 | |
| 492 | """ |
| 493 | input_is_list = isinstance(x, list) |
| 494 | x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent) |
| 495 | time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1) |
| 496 | |
| 497 | if input_img_latents is not None: |
| 498 | input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True) |
| 499 | if input_ids is not None: |
| 500 | condition_embeds = self.llm.embed_tokens(input_ids).clone() |
| 501 | input_img_inx = 0 |
| 502 | for b_inx in input_image_sizes.keys(): |
| 503 | for start_inx, end_inx in input_image_sizes[b_inx]: |
| 504 | condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx] |
| 505 | input_img_inx += 1 |
| 506 | if input_img_latents is not None: |
| 507 | assert input_img_inx == len(input_latents) |
| 508 | |
| 509 | input_emb = torch.cat([condition_embeds, time_token, x], dim=1) |
| 510 | else: |
| 511 | input_emb = torch.cat([time_token, x], dim=1) |
| 512 | 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) |
| 513 | output, past_key_values = output.last_hidden_state, output.past_key_values |
| 514 | if input_is_list: |
| 515 | image_embedding = output[:, -max(num_tokens):] |
| 516 | time_emb = self.t_embedder(timestep, dtype=x.dtype) |
| 517 | x = self.final_layer(image_embedding, time_emb) |
| 518 | latents = [] |
| 519 | for i in range(x.size(0)): |
| 520 | latent = x[i:i+1, :num_tokens[i]] |
| 521 | latent = self.unpatchify(latent, shapes[i][0], shapes[i][1]) |
| 522 | latents.append(latent) |
| 523 | else: |
| 524 | image_embedding = output[:, -num_tokens:] |
| 525 | time_emb = self.t_embedder(timestep, dtype=x.dtype) |
| 526 | x = self.final_layer(image_embedding, time_emb) |
| 527 | latents = self.unpatchify(x, shapes[0], shapes[1]) |
| 528 | |
| 529 | if return_past_key_values: |
| 530 | return latents, past_key_values |
| 531 | return latents |
| 532 | |
| 533 | @torch.no_grad() |
| 534 | def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, offload_model): |
no test coverage detected