Generates encoder outputs from prompt and image input.
(self, config)
| 514 | Generates encoder outputs from prompt and image input. |
| 515 | """ |
| 516 | seed_all(config["seed"]) |
| 517 | self.logger.info("Starting processing in EncoderService...") |
| 518 | room = int(config.get("data_bootstrap_room", 0)) |
| 519 | encoder_metrics = config.setdefault("request_metrics", {}).setdefault("stages", {}).setdefault("encoder", {}) |
| 520 | encoder_metrics["compute_start_ts"] = time.time() |
| 521 | |
| 522 | room_buffers = self._rdma_buffers.get(room) |
| 523 | sender = self.data_sender.get(room) |
| 524 | |
| 525 | prompt = config.get("prompt") |
| 526 | if prompt is None: |
| 527 | raise ValueError("prompt is required in config.") |
| 528 | |
| 529 | # 1. Text Encoding |
| 530 | text_len = config.get("text_len", 512) |
| 531 | |
| 532 | context = self.text_encoder.infer([prompt]) |
| 533 | context = torch.stack([torch.cat([u, u.new_zeros(text_len - u.size(0), u.size(1))]) for u in context]) |
| 534 | |
| 535 | if config.get("enable_cfg", False): |
| 536 | context_null = self.text_encoder.infer([config.get("negative_prompt") or ""]) |
| 537 | context_null = torch.stack([torch.cat([u, u.new_zeros(text_len - u.size(0), u.size(1))]) for u in context_null]) |
| 538 | else: |
| 539 | context_null = None |
| 540 | |
| 541 | text_encoder_output = { |
| 542 | "context": context, |
| 543 | "context_null": context_null, |
| 544 | } |
| 545 | |
| 546 | task = config.get("task") |
| 547 | clip_encoder_out = None |
| 548 | |
| 549 | if task == "t2v": |
| 550 | latent_h = config["size"][0] // config["vae_stride"][1] |
| 551 | latent_w = config["size"][1] // config["vae_stride"][2] |
| 552 | latent_shape = [ |
| 553 | config.get("num_channels_latents", 16), |
| 554 | (config["num_frames"] - 1) // config["vae_stride"][0] + 1, |
| 555 | latent_h, |
| 556 | latent_w, |
| 557 | ] |
| 558 | image_encoder_output = None |
| 559 | elif task == "i2v": |
| 560 | image_path = config.get("image_path") |
| 561 | if image_path is None: |
| 562 | raise ValueError("image_path is required for i2v task.") |
| 563 | |
| 564 | # 2. Image Encoding + VAE Encoding |
| 565 | img, _ = read_image_input(image_path) |
| 566 | |
| 567 | if self.image_encoder is not None: |
| 568 | # Assuming image_encoder.visual handles list of images |
| 569 | clip_encoder_out = self.image_encoder.visual([img]).squeeze(0).to(GET_DTYPE()) |
| 570 | |
| 571 | if self.vae_encoder is None: |
| 572 | raise RuntimeError("VAE encoder is required but was not loaded.") |
| 573 |
no test coverage detected