Views
No views yet
1import torch
2from PIL import Image
3from transformers import AutoProcessor, AutoModelForCausalLM
4from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
5from transformers.image_transforms import resize, to_channel_dimension_format
6from utils import TreeBuilder
7
8
9def convert_to_rgb(image):
10 if image.mode == "RGB":
11 return image
12
13 image_rgba = image.convert("RGBA")
14 background = Image.new("RGBA", image_rgba.size, (255, 255, 255))
15 alpha_composite = Image.alpha_composite(background, image_rgba)
16 alpha_composite = alpha_composite.convert("RGB")
17 return alpha_composite
18
19
20def inference_vlm_websight(image_path, html_path):
21
22 def custom_transform(x):
23 x = convert_to_rgb(x)
24 x = to_numpy_array(x)
25 x = resize(x, (960, 960), resample=PILImageResampling.BILINEAR)
26 x = processor.image_processor.rescale(x, scale=1 / 255)
27 x = processor.image_processor.normalize(
28 x,
29 mean=processor.image_processor.image_mean,
30 std=processor.image_processor.image_std
31 )
32 x = to_channel_dimension_format(x, ChannelDimension.FIRST)
33 x = torch.tensor(x)
34 return x
35
36 model_dir = "lt-asset/Waffle_VLM_WebSight"
37 processor = AutoProcessor.from_pretrained(model_dir)
38 model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
39
40 assert model.config.web_attention_range == 2, "Waffle_VLM_WebSight is trained with hierarchical attention applied to 2 / 8 heads"
41 # use 2/8 = 1/4 attention heads for hierarchical attention (as described in paper)
42 model.eval()
43
44 image_seq_len = model.config.perceiver_config.resampler_n_latents
45 BOS_TOKEN = processor.tokenizer.bos_token
46 BAD_WORDS_IDS = processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
47
48 image = Image.open(image_path)
49 inputs = processor.tokenizer(
50 f"{BOS_TOKEN}<fake_token_around_image>{'<image>' * image_seq_len}<fake_token_around_image>",
51 return_tensors="pt",
52 add_special_tokens=False,
53 )
54 inputs["pixel_values"] = processor.image_processor([image], transform=custom_transform).to(dtype=torch.bfloat16)
55 inputs_for_generation = {k: v.cuda() for k, v in inputs.items()}
56 inputs_for_generation["web_attention_mask"] = None
57 inputs_for_generation["html_tree"] = TreeBuilder(processor.tokenizer)
58 inputs_for_generation["html_tree"].web_attention_mask = inputs_for_generation["web_attention_mask"]
59
60 generated_ids = model.generate(
61 **inputs_for_generation, bad_words_ids=BAD_WORDS_IDS, max_length=2048,
62 num_return_sequences=1, do_sample=False
63 )
64 generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
65
66 with open(html_path, 'w') as wp:
67 wp.write(generated_text)
68
69
70if __name__ == '__main__':
71 inference_vlm_websight('examples/test-495.png', 'examples/example-495.html')
@misc{liang2024wafflemultimodalmodelautomated,
title={WAFFLE: Multi-Modal Model for Automated Front-End Development},
author={Shanchao Liang and Nan Jiang and Shangshu Qian and Lin Tan},
year={2024},
eprint={2410.18362},
archivePrefix={arXiv},
primaryClass={cs.SE},
url={https://arxiv.org/abs/2410.18362},
}