Views
No views yet


1from PIL import Image
2import torch
3from transformers import AutoModel, AutoImageProcessor
4
5MODEL_HUB = "ILLUME-MLLM/dualvitok"
6
7model = AutoModel.from_pretrained(MODEL_HUB, trust_remote_code=True).eval().cuda()
8processor = AutoImageProcessor.from_pretrained(MODEL_HUB, trust_remote_code=True)
9
10# load the diffusion decoder.
11# diffusion_decoder = model.build_sdxl_decoder('ILLUME-MLLM/dualvitok-sdxl-decoder')
12
13# TODO: you need to modify the path here
14IMAGE_PATH = "YOUR_IMAGE_PATH"
15
16image = Image.open(IMAGE_PATH)
17
18image = processor(image, return_tensors="pt")["pixel_values"]
19image = image.cuda()
20
21with torch.no_grad():
22 (quant_semantic, diff_semantic, indices_semantic, _), \
23 (quant_pixel, diff_pixel, indices_pixel) = model.encode(image)
24
25 recon = model.decode(quant_semantic, quant_pixel)
26
27 # decode from the codes.
28 # recon = model.decode_code(indices_semantic, indices_pixel)
29
30 print(recon.shape)
31 recon_image = processor.postprocess(recon)["pixel_values"][0]
32 recon_image.save("recon_image.png")
33
34 # diffusion decoder only support 11 resolution. Check here `diffusion_decoder.resolution_group`.
35 # diffusion_recon = diffusion_decoder(# use vq_indices or vq_embeds
36 # vq_indices=(indices_semantic, indices_pixel),
37 # vq_embeds=(quant_semantic, quant_pixel),
38 # height = height * 2,
39 # width = width * 2,
40 # num_inference_steps = 50,
41 # guidance_scale = 1.5,)
42 # diffusion_recon.images[0].save("diffusion_recon_image.png")