Views
No views yet


1import os
2import os.path as osp
3
4from PIL import Image
5import torch
6from transformers import AutoModel, AutoImageProcessor
7
8MODEL_HUB = "BAAI/Emu3-VisionTokenizer"
9
10model = AutoModel.from_pretrained(MODEL_HUB, trust_remote_code=True).eval().cuda()
11processor = AutoImageProcessor.from_pretrained(MODEL_HUB, trust_remote_code=True)
12
13# TODO: you need to modify the path here
14VIDEO_FRAMES_PATH = "YOUR_VIDEO_FRAMES_PATH"
15
16video = os.listdir(VIDEO_FRAMES_PATH)
17video.sort()
18video = [Image.open(osp.join(VIDEO_FRAMES_PATH, v)) for v in video]
19
20images = processor(video, return_tensors="pt")["pixel_values"]
21images = images.unsqueeze(0).cuda()
22
23# image autoencode
24image = images[:, 0]
25print(image.shape)
26with torch.no_grad():
27 # encode
28 codes = model.encode(image)
29 # decode
30 recon = model.decode(codes)
31
32recon = recon.view(-1, *recon.shape[2:])
33recon_image = processor.postprocess(recon)["pixel_values"][0]
34recon_image.save("recon_image.png")
35
36# video autoencode
37images = images.view(
38 -1,
39 model.config.temporal_downsample_factor,
40 *images.shape[2:],
41)
42
43print(images.shape)
44with torch.no_grad():
45 # encode
46 codes = model.encode(images)
47 # decode
48 recon = model.decode(codes)
49
50recon = recon.view(-1, *recon.shape[2:])
51recon_images = processor.postprocess(recon)["pixel_values"]
52for idx, im in enumerate(recon_images):
53 im.save(f"recon_video_{idx}.png")
54