Views
No views yet
ChatGen-Base-2B is a MLLM finetuned from InternVL-2B. By taking as input a system prompt, and freestyle user query,
the model generates suitable prompts, appropriate models, and specific arguments.ChatGen-Base-2B, first install the necessary dependencies:pip install transformers1import numpy as np
2import torch
3import torchvision.transforms as T
4from PIL import Image
5from torchvision.transforms.functional import InterpolationMode
6from transformers import AutoModel, AutoTokenizer
7
8IMAGENET_MEAN = (0.485, 0.456, 0.406)
9IMAGENET_STD = (0.229, 0.224, 0.225)
10
11def build_transform(input_size):
12 MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
13 transform = T.Compose([
14 T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
15 T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
16 T.ToTensor(),
17 T.Normalize(mean=MEAN, std=STD)
18 ])
19 return transform
20
21def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
22 best_ratio_diff = float('inf')
23 best_ratio = (1, 1)
24 area = width * height
25 for ratio in target_ratios:
26 target_aspect_ratio = ratio[0] / ratio[1]
27 ratio_diff = abs(aspect_ratio - target_aspect_ratio)
28 if ratio_diff < best_ratio_diff:
29 best_ratio_diff = ratio_diff
30 best_ratio = ratio
31 elif ratio_diff == best_ratio_diff:
32 if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
33 best_ratio = ratio
34 return best_ratio
35
36def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
37 orig_width, orig_height = image.size
38 aspect_ratio = orig_width / orig_height
39
40 # calculate the existing image aspect ratio
41 target_ratios = set(
42 (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
43 i * j <= max_num and i * j >= min_num)
44 target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
45
46 # find the closest aspect ratio to the target
47 target_aspect_ratio = find_closest_aspect_ratio(
48 aspect_ratio, target_ratios, orig_width, orig_height, image_size)
49
50 # calculate the target width and height
51 target_width = image_size * target_aspect_ratio[0]
52 target_height = image_size * target_aspect_ratio[1]
53 blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
54
55 # resize the image
56 resized_img = image.resize((target_width, target_height))
57 processed_images = []
58 for i in range(blocks):
59 box = (
60 (i % (target_width // image_size)) * image_size,
61 (i // (target_width // image_size)) * image_size,
62 ((i % (target_width // image_size)) + 1) * image_size,
63 ((i // (target_width // image_size)) + 1) * image_size
64 )
65 # split the image
66 split_img = resized_img.crop(box)
67 processed_images.append(split_img)
68 assert len(processed_images) == blocks
69 if use_thumbnail and len(processed_images) != 1:
70 thumbnail_img = image.resize((image_size, image_size))
71 processed_images.append(thumbnail_img)
72 return processed_images
73
74def load_image(image_file, input_size=448, max_num=12):
75 image = Image.open(image_file).convert('RGB')
76 transform = build_transform(input_size=input_size)
77 images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
78 pixel_values = [transform(image) for image in images]
79 pixel_values = torch.stack(pixel_values)
80 return pixel_values
81
82# If you want to load a model using multiple GPUs, please refer to the `Multiple GPUs` section.
83path = 'ChengyouJia/ChatGen-Base-2B'
84model = AutoModel.from_pretrained(
85 path,
86 torch_dtype=torch.bfloat16,
87 low_cpu_mem_usage=True,
88 trust_remote_code=True).eval().cuda()
89tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
90
91sys_singlemodal = """
92You are a user requirements translation expert. I have a freestyle prompt written by a non professional user for text-to-image tasks. Please convert the content of this freestyle prompt into professional prompt and professional negativePrompt, and provide the model and its parameters that are most suitable for the user's text-to-image task.
93Here is the content I need you to convert:
94"""
95
96sys_multimodal = """
97You are a user requirements translation expert. I have a freestyle prompt written by a non professional user for text-to-image tasks.
98Additionally, a general user provide several reference images, indicating that they want the final generated image to have a style similar to those images. You should combine the reference images to convert the content of the freestyle prompt into professional prompt and professional negativePrompt, and provide the model and its parameters that are most suitable for the user's text-to-image task.
99Here are the reference images and content I need you to convert:
100"""
101
102# set the max number of tiles in `max_num`
103pixel_values = None
104<!-- pixel_values = load_image(<image_path>, max_num=6).to(torch.bfloat16).cuda() -->
105generation_config = dict(max_new_tokens=1024, do_sample=True)
106
107question = "Whip up a cool sci-fi robot girl, colorful and detailed from waist up, y'know?"
108
109input = sys_singlemodal + question
110response, history = model.chat(tokenizer, None, input, generation_config, history=None, return_history=True)
111print(f'User: {question}\nAssistant: {response}')
## Citation
If you find this repository helpful, feel free to cite our paper:
```bibtex
@article{jia2024chatgen,
title={ChatGen: Automatic Text-to-Image Generation From FreeStyle Chatting},
author={Jia, Chengyou and Xia, Changliang and Dang, Zhuohang and Wu, Weijia and Qian, Hangwei and Luo, Minnan},
journal={arXiv preprint arXiv:2411.17176},
year={2024}
}