| 38 | |
| 39 | |
| 40 | class Sampler: |
| 41 | def __init__(self): |
| 42 | self.mesh = VideoLLaMAConfig.get_jax_mesh(FLAGS.mesh_dim) |
| 43 | self.vqgan = VQGAN(FLAGS.vqgan_checkpoint, replicate=False) |
| 44 | self.prefix_tokenizer = AutoTokenizer.from_pretrained(FLAGS.tokenizer, truncation_side='left', padding_side='left') |
| 45 | self.tokenizer = AutoTokenizer.from_pretrained(FLAGS.tokenizer) |
| 46 | self.n_tokens_per_frame = 257 |
| 47 | self.min_buffer_size = 256 |
| 48 | self.sharded_rng = next_rng() |
| 49 | self._load_model() |
| 50 | |
| 51 | @property |
| 52 | def block_size(self): |
| 53 | return max(self.config.scan_query_chunk_size, self.config.scan_key_chunk_size) * self.mesh.shape['sp'] |
| 54 | |
| 55 | @property |
| 56 | def data_dim(self): |
| 57 | return self.mesh.shape['dp'] * self.mesh.shape['fsdp'] |
| 58 | |
| 59 | def _process_frame(self, image, size): |
| 60 | width, height = image.size |
| 61 | if width < height: |
| 62 | new_width = size |
| 63 | new_height = int(size * height / width) |
| 64 | else: |
| 65 | new_height = size |
| 66 | new_width = int(size * width / height) |
| 67 | image = image.resize((new_width, new_height)) |
| 68 | |
| 69 | left = (new_width - size) / 2 |
| 70 | top = (new_height - size) / 2 |
| 71 | right = (new_width + size) / 2 |
| 72 | bottom = (new_height + size) / 2 |
| 73 | image = image.crop((left, top, right, bottom)) |
| 74 | return np.array(image, dtype=np.float32) / 127.5 - 1 |
| 75 | |
| 76 | def _read_process_vision(self, path, max_n_frames): |
| 77 | f = open_file(path, 'rb') |
| 78 | if path.endswith('.png') or path.endswith('.jpg'): |
| 79 | image = Image.open(f).convert('RGB') |
| 80 | vision = self._process_frame(image, 256)[None] |
| 81 | else: |
| 82 | vr = decord.VideoReader(f, ctx=decord.cpu(0)) |
| 83 | duration = len(vr) |
| 84 | if duration <= max_n_frames: |
| 85 | frame_id_list = list(range(duration)) |
| 86 | else: |
| 87 | frame_id_list = np.linspace(0, duration - 1, max_n_frames, dtype=int).tolist() |
| 88 | video = vr.get_batch(frame_id_list).asnumpy() |
| 89 | vision = np.stack([self._process_frame(Image.fromarray(frame), 256) for frame in video]) |
| 90 | |
| 91 | B = 1 |
| 92 | encodings = [] |
| 93 | for i in range(0, len(vision), 1): |
| 94 | v = vision[i:i + B] |
| 95 | if len(v) % B == 0: |
| 96 | n_pad = 0 |
| 97 | else: |