Method generate_once

evaluation/vqlm_demo/inference.py:103–159

(self, input_images, n_new_frames, temperature=1.0, top_p=1.0)

    @torch.no_grad()
    def generate_once(self, input_images, n_new_frames, temperature=1.0, top_p=1.0):
        assert type(input_images) == np.ndarray
        with self.lock:
            input_images = np.array(input_images, dtype=np.float32)
            input_images = torch.tensor(
                einops.rearrange(input_images, 'b h w c -> b c h w')
            ).to(self.torch_device)

            print('here:', type(input_images))

            # old tokenizer
            # input_ids = self.tokenizer.tokenize(input_images).view(1, -1)

            # new tokenizer
            _, input_ids = self.tokenizer.encode(input_images)
            input_ids = input_ids.view(1, -1)


            input_ids = input_ids[:, -(self.context_frames - 1) * 256:]

            new_tokens = []
            current_context_frames = input_ids.shape[1] // 256
            fisrt_generation_left = self.context_frames - current_context_frames
            first_new_frames = min(fisrt_generation_left, n_new_frames)
            input_ids = self.model.generate(
                input_ids=input_ids,
                attention_mask=torch.ones_like(input_ids),
                pad_token_id=8192,
                max_new_tokens=256 * first_new_frames,
                do_sample=True,
                top_p=top_p,
                temperature=temperature,
                suppress_tokens=list(range(8192, self.model.vocab_size)),
            )
            new_tokens.append(input_ids[:, -256 * first_new_frames:])
            input_ids = input_ids[:, -(self.context_frames - 1) * 256:]

            for _ in range(max(0, n_new_frames - first_new_frames)):
                input_ids = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=torch.ones_like(input_ids),
                    pad_token_id=8192,
                    max_new_tokens=256,
                    do_sample=True,
                    top_p=top_p,
                    temperature=temperature,
                    suppress_tokens=list(range(8192, self.model.vocab_size)),
                )
                new_tokens.append(input_ids[:, -256:])
                input_ids = input_ids[:, -(self.context_frames - 1) * 256:]

            new_tokens = torch.cat(new_tokens, dim=1).view(-1, 256)
            new_images = einops.rearrange(
                torch.clamp(self.tokenizer.decode_code(new_tokens), 0.0, 1.0),
                'b c h w -> b h w c'
            ).detach().cpu().numpy()
        return new_images