Method cropped_pos_embed

tensorrt_llm/layers/embedding.py:358–390

Crops positional embeddings for SD3 compatibility.

(self, height, width)

                f"Unsupported pos_embed_type: {self.pos_embed_type}")

    def cropped_pos_embed(self, height, width):
        """Crops positional embeddings for SD3 compatibility."""
        if self.pos_embed_max_size is None:
            raise ValueError("`pos_embed_max_size` must be set for cropping.")

        height = height // self.patch_size
        width = width // self.patch_size
        if height > self.pos_embed_max_size:
            raise ValueError(
                f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
            )
        if width > self.pos_embed_max_size:
            raise ValueError(
                f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
            )

        top = (self.pos_embed_max_size - height) // 2
        left = (self.pos_embed_max_size - width) // 2
        spatial_pos_embed = identity(self.pos_embed.value).view(
            [1, self.pos_embed_max_size, self.pos_embed_max_size, -1])
        spatial_pos_embed = slice(spatial_pos_embed,
                                  starts=[0, top, left, 0],
                                  sizes=concat([
                                      shape(spatial_pos_embed, 0), height,
                                      width,
                                      shape(spatial_pos_embed, 3)
                                  ]))
        spatial_pos_embed = spatial_pos_embed.view(
            concat(
                [1, -1,
                 shape(spatial_pos_embed,
                       spatial_pos_embed.ndim() - 1)]))
        return spatial_pos_embed

    def forward(self, latent):
        # [TODO] to support height and width for runtime