Method prepare_inputs

tensorrt_llm/models/stdit/model.py:1280–1441

@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the ranges of the dimensions of when using TRT dynamic shapes. @return: a list contains values which can be fed into the self.forward()

(self, max_batch_size, **kwargs)

        return output

    def prepare_inputs(self, max_batch_size, **kwargs):
        '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the
            ranges of the dimensions of when using TRT dynamic shapes.

            @return: a list contains values which can be fed into the self.forward()
        '''

        mapping = self.config.mapping
        if mapping.tp_size > 1:
            current_all_reduce_helper().set_workspace_tensor(mapping, 1)

        def stdit_default_batch_range(max_batch_size):
            return [max_batch_size, max_batch_size, max_batch_size]

        default_range = stdit_default_batch_range
        # [NOTE] For now only static batch size is supported, so we run the model with max_batch_size.
        batch_size = max_batch_size

        x = Tensor(name='x',
                   dtype=self.dtype,
                   shape=[batch_size, self.in_channels, *self.latent_size],
                   dim_range=OrderedDict([
                       ('batch_size', [default_range(max_batch_size)]),
                       ('in_channels', [[self.in_channels] * 3]),
                       ('latent_frames', [[self.latent_size[0]] * 3]),
                       ('latent_height', [[self.latent_size[1]] * 3]),
                       ('latent_width', [[self.latent_size[2]] * 3]),
                   ]))
        timestep = Tensor(name='timestep',
                          dtype=self.dtype,
                          shape=[batch_size],
                          dim_range=OrderedDict([
                              ('batch_size', [default_range(max_batch_size)]),
                          ]))
        y = Tensor(
            name='y',
            dtype=self.dtype,
            shape=[batch_size, 1, self.model_max_length, self.caption_channels],
            dim_range=OrderedDict([
                ('batch_size', [default_range(max_batch_size)]),
                ('mask_batch_size', [[1, 1, 1]]),
                ('num_tokens', [[self.model_max_length] * 3]),
                ('caption_channels', [[self.caption_channels] * 3]),
            ]))
        mask = Tensor(name='mask',
                      dtype=trt.int32,
                      shape=[1, self.model_max_length],
                      dim_range=OrderedDict([
                          ('mask_batch_size', [[1, 1, 1]]),
                          ('num_tokens', [[self.model_max_length] * 3]),
                      ]))
        x_mask = Tensor(name='x_mask',
                        dtype=tensorrt_llm.str_dtype_to_trt('bool'),
                        shape=[batch_size, self.latent_size[0]],
                        dim_range=OrderedDict([
                            ('batch_size', [default_range(max_batch_size)]),
                            ('latent_frames', [[self.latent_size[0]] * 3]),
                        ]))