tttoaster commited on
Commit
1504887
·
verified ·
1 Parent(s): b1eafeb

Delete src/demo

Browse files
src/demo/__pycache__/conversation.cpython-311.pyc DELETED
Binary file (8.21 kB)
 
src/demo/__pycache__/conversation.cpython-38.pyc DELETED
Binary file (4.43 kB)
 
src/demo/__pycache__/utils.cpython-311.pyc DELETED
Binary file (4.39 kB)
 
src/demo/__pycache__/utils.cpython-38.pyc DELETED
Binary file (2.32 kB)
 
src/demo/conversation.py DELETED
@@ -1,182 +0,0 @@
1
- import dataclasses
2
- from enum import auto, Enum
3
- from typing import List, Tuple
4
-
5
- import io
6
- import base64
7
- import os
8
- from PIL import Image
9
- import copy
10
-
11
- IMG_FLAG = '<image>'
12
-
13
-
14
- class SeparatorStyle(Enum):
15
- """Different separator style."""
16
- SINGLE = auto()
17
- TWO = auto()
18
- MPT = auto()
19
- PLAIN = auto()
20
- LLAMA_2 = auto()
21
-
22
-
23
- def decode_image(encoded_image: str) -> Image:
24
- decoded_bytes = base64.b64decode(encoded_image.encode('utf-8'))
25
- buffer = io.BytesIO(decoded_bytes)
26
- image = Image.open(buffer)
27
- return image
28
-
29
-
30
- def encode_image(image: Image.Image, format: str = 'PNG') -> str:
31
- with io.BytesIO() as buffer:
32
- image.save(buffer, format=format)
33
- encoded_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
34
- return encoded_image
35
-
36
-
37
- @dataclasses.dataclass
38
- class Conversation:
39
- """A class that keeps all conversation history."""
40
- system: str
41
- roles: List[str]
42
- messages: List[dict] # multi-turn -> user & assistant -> {'images': [PIL.Image,], 'text': str}
43
- offset: int
44
- sep_style: SeparatorStyle = SeparatorStyle.SINGLE
45
- sep: str = "###"
46
- sep2: str = None
47
- version: str = "Unknown"
48
-
49
- skip_next: bool = False
50
-
51
- def get_prompt(self):
52
- messages = copy.deepcopy(self.messages)
53
- if self.sep_style == SeparatorStyle.SINGLE:
54
- if self.system is None or self.system == '':
55
- text = ''
56
- else:
57
- text = self.system + self.sep
58
- images = []
59
- for message in messages:
60
- text += message['role'] + ": " + message['message']['text'] + self.sep
61
- for image_path in message['message']['images']:
62
- image = Image.open(image_path).resize((256, 256))
63
- image_base64 = encode_image(image)
64
- images.append(image_base64)
65
-
66
- text += self.roles[1] + ":"
67
- elif self.sep_style == SeparatorStyle.LLAMA_2:
68
- b_token = "[INST] "
69
- e_token = " [/INST]"
70
- if self.system is None or self.system == '':
71
- text = ''
72
- else:
73
- text = f"<<SYS>>\n{self.system}\n<</SYS>>\n\n"
74
- images = []
75
- for idx, message in enumerate(messages):
76
- # text += message['role'] + ": " + message['message']['text'] + self.sep
77
- if idx % 2 == 0:
78
- text += b_token + message['message']['text'] + e_token + self.sep
79
- else:
80
- text += message['message']['text'] + self.sep
81
-
82
- for image_path in message['message']['images']:
83
- image = Image.open(image_path)
84
- image_base64 = encode_image(image)
85
- images.append(image_base64)
86
- else:
87
- raise NotImplementedError
88
-
89
- return {'text': text, 'images': images}
90
-
91
- # def update_image_ids(self, images_ids):
92
- # image_count = 0
93
- # for message in self.messages:
94
- # for idx in range(len(message['message']['images_ids'])):
95
- # if message['message']["images_ids"][idx] is None:
96
- # message['message']["images_ids"][idx] = images_ids[image_count]
97
- # image_count += 1
98
-
99
- # assert len(images_ids) == image_count, print(len(images_ids), image_count)
100
-
101
- def append_message(self, role, message):
102
- self.messages.append([role, message])
103
-
104
- def to_gradio_chatbot(self):
105
- dialog = []
106
- for i, single_turn in enumerate(self.messages[self.offset:]):
107
- single_turn = single_turn['message']
108
- text_list = single_turn['text'].split(IMG_FLAG)
109
- assert len(text_list) == len(single_turn['images']) + 1, print(text_list, len(single_turn['images']))
110
- message = ''
111
- for image_idx in range(len(single_turn['images'])):
112
- # image = single_turn['images'][image_idx]
113
- # image_base64 = encode_image(image)
114
- # image_str = f'<img src="data:image/png;base64,{image_base64}" alt="user upload image" />'
115
- image_path = single_turn['images'][image_idx]
116
- if image_path == '':
117
- message += text_list[image_idx] + '<corrupt_image>'
118
- else:
119
- message += text_list[image_idx] + f'![](file={image_path})'
120
- message += text_list[-1]
121
-
122
- if i % 2 == 0:
123
- dialog.append([message, None])
124
- else:
125
- dialog[-1][-1] = message
126
-
127
- return dialog
128
-
129
- def copy(self):
130
- return Conversation(system=self.system,
131
- roles=self.roles,
132
- messages=copy.deepcopy(self.messages),
133
- offset=self.offset,
134
- sep_style=self.sep_style,
135
- sep=self.sep,
136
- sep2=self.sep2,
137
- version=self.version)
138
-
139
- def dict(self):
140
- messages = copy.deepcopy(self.messages)
141
- for message in messages:
142
- for i in range(len(message['message']['images'])):
143
- message['message']['images'][i] = os.path.basename(message['message']['images'][i])
144
- return {
145
- "system": self.system,
146
- "roles": self.roles,
147
- "messages": messages,
148
- "offset": self.offset,
149
- "sep": self.sep,
150
- "sep2": self.sep2,
151
- }
152
-
153
-
154
- conv_seed_vicuna = Conversation(
155
- system="",
156
- roles=("USER", "ASSISTANT"),
157
- version="v2",
158
- messages=[],
159
- offset=0,
160
- sep_style=SeparatorStyle.SINGLE,
161
- sep='\n',
162
- )
163
-
164
- conv_seed_vicuna_system = Conversation(
165
- system="A chat between a curious user and an artificial intelligence assistant. ",
166
- roles=("USER", "ASSISTANT"),
167
- version="v2",
168
- messages=[],
169
- offset=0,
170
- sep_style=SeparatorStyle.SINGLE,
171
- sep='\n',
172
- )
173
-
174
- conv_seed_llama2 = Conversation(
175
- system="",
176
- roles=("[INST]", "[/INST]"),
177
- version="v2",
178
- messages=[],
179
- offset=0,
180
- sep_style=SeparatorStyle.LLAMA_2,
181
- sep='\n',
182
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/demo/seed_llama_flask.py DELETED
@@ -1,379 +0,0 @@
1
- import hydra
2
- import pyrootutils
3
- import torch
4
- import re
5
- import time
6
- from omegaconf import OmegaConf
7
- from flask import Flask, request
8
- from typing import Optional
9
- import transformers
10
- from dataclasses import dataclass, field
11
- import io
12
- import base64
13
- from PIL import Image
14
- import numpy as np
15
- import cv2
16
- from diffusers import AutoencoderKL, UNet2DConditionModel, EulerDiscreteScheduler
17
-
18
-
19
- pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
20
-
21
- from src.data.any_res import process_anyres_image
22
-
23
- BOI_TOKEN = '<img>'
24
- BOP_TOKEN = '<patch>'
25
- EOI_TOKEN = '</img>'
26
- EOP_TOKEN = '</patch>'
27
- IMG_TOKEN = '<img_{:05d}>'
28
-
29
- IMG_FLAG = '<image>'
30
- num_img_in_tokens = 64
31
- num_img_out_tokens = 64
32
-
33
- resolution_grids = ['1x1', '1x2', '1x3', '1x4', '1x5', '1x6', '1x10', '2x1', '3x1', '4x1', '5x1', '6x1', '10x1', '2x2', '2x3', '3x2', '2x4', '4x2']
34
- base_resolution = 448
35
-
36
- app = Flask(__name__)
37
-
38
-
39
- def decode_image(encoded_image: str) -> Image:
40
- decoded_bytes = base64.b64decode(encoded_image.encode('utf-8'))
41
- buffer = io.BytesIO(decoded_bytes)
42
- image = Image.open(buffer)
43
- return image
44
-
45
-
46
- def encode_image(image: Image.Image, format: str = 'PNG') -> str:
47
- with io.BytesIO() as buffer:
48
- image.save(buffer, format=format)
49
- encoded_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
50
- return encoded_image
51
-
52
-
53
- @dataclass
54
- class Arguments:
55
- image_transform: Optional[str] = field(default=None, metadata={"help": "config path of image transform"})
56
- tokenizer: Optional[str] = field(default=None, metadata={"help": "config path of tokenizer used to initialize tokenizer"})
57
- llm: Optional[str] = field(default=None, metadata={"help": "config path of llm"})
58
- visual_encoder: Optional[str] = field(default=None, metadata={"help": "config path of visual encoder"})
59
- sd_adapter: Optional[str] = field(default=None, metadata={"help": "config path of sd adapter"})
60
- agent: Optional[str] = field(default=None, metadata={"help": "config path of agent model"})
61
- diffusion_path: Optional[str] = field(default=None, metadata={"help": "diffusion model path"})
62
- has_bbox: Optional[bool] = field(default=False, metadata={"help": "visualize the box"})
63
-
64
- port: Optional[str] = field(default=80, metadata={"help": "network port"})
65
- llm_device: Optional[str] = field(default='cuda:0', metadata={"help": "llm device"})
66
- vit_sd_device: Optional[str] = field(default='cuda:0', metadata={"help": "sd and vit device"})
67
- dtype: Optional[str] = field(default='fp16', metadata={"help": "mix percision"})
68
-
69
- multi_resolution: Optional[bool] = field(default=False, metadata={"help": "multi resolution"})
70
-
71
-
72
- parser = transformers.HfArgumentParser(Arguments)
73
- args, = parser.parse_args_into_dataclasses()
74
-
75
- def extract_box(output_str):
76
- boxes = re.findall('(.*?)<box_end>', output_str)
77
- if len(boxes) >0:
78
- bboxes = [[int(num) for num in re.findall('<loc-(\d+)>', box)] for box in boxes]
79
- else:
80
- bboxes = None
81
-
82
- return bboxes
83
-
84
-
85
- def visualize_bbox(image, bboxes):
86
- img_width, img_height = image.size
87
- image = np.array(image)
88
- image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
89
- for bbox in bboxes:
90
- x_center, y_center, box_width, box_height = bbox
91
-
92
- x_center = x_center / 224 * img_width
93
- y_center = y_center / 224 * img_height
94
-
95
- box_width = box_width /224 * img_width
96
- box_height = box_height / 224 * img_height
97
-
98
- x1 = int(x_center - box_width / 2)
99
- y1 = int(y_center - box_height / 2)
100
- x2 = int(x_center + box_width / 2)
101
- y2 = int(y_center + box_height / 2)
102
-
103
- cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 4)
104
-
105
- image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
106
- image = Image.fromarray(image)
107
-
108
-
109
- return image
110
-
111
-
112
-
113
-
114
- class LLMService:
115
-
116
- def __init__(self, args) -> None:
117
-
118
- self.llm_device = args.llm_device
119
- self.vit_sd_device = args.vit_sd_device
120
-
121
- dtype = args.dtype
122
- if dtype == 'fp16':
123
- self.dtype = torch.float16
124
- elif dtype == 'bf16':
125
- self.dtype = torch.bfloat16
126
- else:
127
- raise ValueError
128
-
129
- image_transform_cfg = OmegaConf.load(args.image_transform)
130
- self.image_transform = hydra.utils.instantiate(image_transform_cfg)
131
-
132
- tokenizer_cfg = OmegaConf.load(args.tokenizer)
133
- self.tokenizer = hydra.utils.instantiate(tokenizer_cfg)
134
-
135
- visual_encoder_cfg = OmegaConf.load(args.visual_encoder)
136
- self.visual_encoder = hydra.utils.instantiate(visual_encoder_cfg)
137
- self.visual_encoder.eval().to(self.vit_sd_device, dtype=self.dtype)
138
- print('Init visual encoder done')
139
-
140
- llm_cfg = OmegaConf.load(args.llm)
141
- llm = hydra.utils.instantiate(llm_cfg, torch_dtype=self.dtype)
142
- print('Init llm done.')
143
-
144
- agent_cfg = OmegaConf.load(args.agent)
145
- self.agent = hydra.utils.instantiate(agent_cfg, llm=llm)
146
-
147
- self.agent.eval().to(self.llm_device, dtype=self.dtype)
148
- print('Init agent mdoel Done')
149
-
150
- noise_scheduler = EulerDiscreteScheduler.from_pretrained(args.diffusion_path, subfolder="scheduler")
151
-
152
- vae = AutoencoderKL.from_pretrained(args.diffusion_path, subfolder="vae").to(self.vit_sd_device, dtype=self.dtype)
153
-
154
- unet = UNet2DConditionModel.from_pretrained(args.diffusion_path, subfolder="unet").to(dtype=self.dtype)
155
-
156
- sd_adapter_cfg = OmegaConf.load(args.sd_adapter)
157
-
158
- self.sd_adapter = hydra.utils.instantiate(sd_adapter_cfg, unet=unet).eval().to(dtype=self.dtype)
159
-
160
- self.sd_adapter.init_pipe(vae=vae,
161
- scheduler=noise_scheduler,
162
- visual_encoder=self.visual_encoder.to("cpu"),
163
- image_transform=self.image_transform,
164
- discrete_model=None,
165
- dtype=self.dtype,
166
- device="cpu")
167
-
168
- print('Init sd adapter pipe done.')
169
-
170
- self.visual_encoder.to(self.vit_sd_device, dtype=self.dtype)
171
-
172
- self.boi_token_id = self.tokenizer.encode(BOI_TOKEN, add_special_tokens=False)[0]
173
- self.eoi_token_id = self.tokenizer.encode(EOI_TOKEN, add_special_tokens=False)[0]
174
-
175
- self.bop_token_id = self.tokenizer.encode(BOP_TOKEN, add_special_tokens=False)[0]
176
- self.eop_token_id = self.tokenizer.encode(EOP_TOKEN, add_special_tokens=False)[0]
177
-
178
- self.multi_resolution = args.multi_resolution
179
- if self.multi_resolution:
180
- self.base_resolution = base_resolution
181
- grid_pinpoints = []
182
- for scale in resolution_grids:
183
- s1, s2 = scale.split('x')
184
- grid_pinpoints.append([int(s1)*base_resolution, int(s2)*base_resolution])
185
- self.grid_pinpoints = grid_pinpoints
186
-
187
-
188
- service = LLMService(args)
189
-
190
-
191
- @app.route('/generate', methods=['GET', 'POST'])
192
- def generate():
193
- with torch.no_grad():
194
- request_info = request.get_json()
195
-
196
- text_list = request_info['text'].split(IMG_FLAG)
197
- image_list = request_info['images']
198
- max_new_tokens = request_info.get('max_new_tokens', 256)
199
- top_p = 0.5
200
- force_boi = request_info.get('force_boi', False)
201
- force_bbox = request_info.get('force_bbox', False)
202
-
203
- assert len(text_list) == len(image_list) + 1
204
-
205
- image_tokens = BOI_TOKEN + ''.join([IMG_TOKEN.format(int(item)) for item in range(num_img_in_tokens)]) + EOI_TOKEN
206
-
207
- input_images = []
208
- if len(image_list) > 0:
209
- image_tensor_list = []
210
- embeds_cmp_mask = []
211
- embeds_gen_mask = []
212
-
213
- if service.multi_resolution:
214
- patch_pos = []
215
- image_patch_length = []
216
- image_size_list = []
217
-
218
- for idx, image_item in enumerate(image_list):
219
- if isinstance(image_item, str):
220
- image = decode_image(image_item)
221
- print('after decode image size:', image.size)
222
- input_images.append(image)
223
-
224
- if service.multi_resolution:
225
- image_size_list.append(image.size)
226
- print('image size:', image.size)
227
- image_tensor, patch_pos_tensor = process_anyres_image(image, service.image_transform, service.grid_pinpoints, service.base_resolution)
228
- image_tensor_list.append(image_tensor)
229
- patch_pos.append(patch_pos_tensor)
230
- image_patch_length.append(image_tensor.shape[0])
231
- print('image_patch_length', image_patch_length)
232
- embeds_cmp_mask.extend([True]*image_tensor.shape[0])
233
- embeds_gen_mask.extend([False]*image_tensor.shape[0])
234
-
235
- else:
236
- image_tensor = service.image_transform(image)
237
-
238
- image_tensor_list.append(image_tensor)
239
- embeds_cmp_mask.append(True)
240
- embeds_gen_mask.append(False)
241
- else:
242
- raise ValueError
243
-
244
- if service.multi_resolution:
245
- pixel_values = torch.cat(image_tensor_list).to(service.vit_sd_device, dtype=service.dtype)
246
- patch_position = torch.cat(patch_pos, dim=0)
247
-
248
- image_tokens_list = []
249
- for patch_length in image_patch_length:
250
- image_tokens = ''
251
- for _ in range(patch_length-1):
252
- image_tokens += BOP_TOKEN + ''.join(IMG_TOKEN.format(int(item)) for item in range(num_img_in_tokens)) + EOP_TOKEN
253
- image_tokens += BOI_TOKEN + ''.join(IMG_TOKEN.format(int(item)) for item in range(num_img_in_tokens)) + EOI_TOKEN
254
- image_tokens_list.append(image_tokens)
255
- else:
256
- pixel_values = torch.stack(image_tensor_list).to(service.vit_sd_device, dtype=service.dtype)
257
-
258
- image_embeds = service.visual_encoder(pixel_values)
259
- image_embeds = image_embeds.to(service.llm_device)
260
-
261
- embeds_cmp_mask = torch.tensor(embeds_cmp_mask, dtype=torch.bool).to(service.llm_device)
262
- embeds_gen_mask = torch.tensor(embeds_gen_mask, dtype=torch.bool).to(service.llm_device)
263
-
264
- else:
265
- image_embeds = None
266
- patch_position = 0
267
- embeds_cmp_mask = None
268
- embeds_gen_mask = None
269
-
270
- if service.multi_resolution:
271
- input_text = ''
272
- for i, c in enumerate(text_list[:-1]):
273
- input_text += c + image_tokens_list[i]
274
- input_text += text_list[-1]
275
-
276
- else:
277
- input_text = image_tokens.join(text_list)
278
-
279
- if force_boi:
280
- input_text = input_text + BOI_TOKEN
281
-
282
- if force_bbox:
283
- input_text = input_text + '[[ <box_start>'
284
- print('input_text:', input_text)
285
- input_ids = service.tokenizer.encode(input_text, add_special_tokens=False)
286
- input_ids = [service.tokenizer.bos_token_id] + input_ids
287
-
288
- input_ids = torch.tensor(input_ids).to(service.llm_device, dtype=torch.long)
289
- ids_cmp_mask = torch.zeros_like(input_ids, dtype=torch.bool).to(service.llm_device)
290
- ids_gen_mask = torch.zeros_like(input_ids, dtype=torch.bool).to(service.llm_device)
291
-
292
- if service.multi_resolution:
293
- boi_indices = torch.where(torch.logical_or(input_ids == service.boi_token_id, input_ids == service.bop_token_id))[0].tolist()
294
- eoi_indices = torch.where(torch.logical_or(input_ids == service.eoi_token_id, input_ids == service.eop_token_id))[0].tolist()
295
-
296
- else:
297
-
298
- boi_indices = torch.where(input_ids == service.boi_token_id)[0].tolist()
299
- eoi_indices = torch.where(input_ids == service.eoi_token_id)[0].tolist()
300
-
301
- for boi_idx, eoi_idx in zip(boi_indices, eoi_indices):
302
- ids_cmp_mask[boi_idx + 1:eoi_idx] = True
303
-
304
- input_ids = input_ids.unsqueeze(0)
305
- ids_cmp_mask = ids_cmp_mask.unsqueeze(0)
306
- ids_gen_mask = ids_gen_mask.unsqueeze(0)
307
-
308
- error_msg = []
309
-
310
- if service.multi_resolution:
311
- output = service.agent.generate(
312
- tokenizer=service.tokenizer,
313
- input_ids=input_ids,
314
- image_embeds=image_embeds,
315
- patch_positions=patch_position,
316
- embeds_cmp_mask=embeds_cmp_mask,
317
- ids_cmp_mask=ids_cmp_mask,
318
- num_img_gen_tokens=num_img_out_tokens,
319
- max_new_tokens=max_new_tokens,
320
- dtype=service.dtype,
321
- device=service.llm_device,
322
- top_p=top_p,
323
- )
324
- else:
325
- output = service.agent.generate(
326
- tokenizer=service.tokenizer,
327
- input_ids=input_ids,
328
- image_embeds=image_embeds,
329
- embeds_cmp_mask=embeds_cmp_mask,
330
- ids_cmp_mask=ids_cmp_mask,
331
- num_img_gen_tokens=num_img_out_tokens,
332
- max_new_tokens=max_new_tokens,
333
- dtype=service.dtype,
334
- device=service.llm_device,
335
- top_p=top_p,
336
- )
337
-
338
- gen_imgs_base64_list = []
339
- generated_text = output['text']
340
- generated_text = generated_text.replace(EOI_TOKEN, IMG_FLAG).replace(service.tokenizer.eos_token, '')
341
-
342
- if output['has_img_output']:
343
- print('loading visual encoder and llm to CPU, and sd to GPU')
344
- a = time.time()
345
- service.agent = service.agent.to("cpu")
346
- service.sd_adapter = service.sd_adapter.to(service.vit_sd_device, dtype=service.dtype)
347
- print("Loading finished: ", time.time() - a)
348
-
349
- img_gen_feat = output['img_gen_feat'].to(service.vit_sd_device, dtype=service.dtype)
350
-
351
- for img_idx in range(output['num_gen_imgs']):
352
- img_feat = img_gen_feat[img_idx:img_idx + 1]
353
- generated_image = service.sd_adapter.generate(image_embeds=img_feat, num_inference_steps=50)[0]
354
- image_base64 = encode_image(generated_image)
355
- gen_imgs_base64_list.append(image_base64)
356
-
357
- print('loading visual encoder and llm to GPU, and sd to CPU')
358
- a = time.time()
359
- service.sd_adapter = service.sd_adapter.to("cpu")
360
- service.visual_encoder = service.visual_encoder.to(service.vit_sd_device, dtype=service.dtype)
361
- service.agent = service.agent.to(service.vit_sd_device, dtype=service.dtype)
362
- print("Loading finished: ", time.time() - a)
363
-
364
- if args.has_bbox:
365
- bboxes = extract_box(generated_text)
366
-
367
- if bboxes is not None and len(input_images) > 0:
368
- image_viz = visualize_bbox(input_images[0], bboxes)
369
- image_base64 = encode_image(image_viz)
370
- gen_imgs_base64_list.append(image_base64)
371
- generated_text = re.sub(r'\[\[ <box_start>.*?<box_end>.*?\]\]', 'the green bounding box', generated_text)
372
- generated_text += IMG_FLAG
373
- print(input_text + generated_text)
374
-
375
- return {'text': generated_text, 'images': gen_imgs_base64_list, 'error_msg': error_msg}
376
-
377
-
378
- if __name__ == '__main__':
379
- app.run(host='0.0.0.0', port=args.port)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/demo/seed_llama_gradio.py DELETED
@@ -1,465 +0,0 @@
1
- import os
2
- import numpy as np
3
- import datetime
4
- import json
5
- from typing import Optional
6
- import transformers
7
- from dataclasses import dataclass, field
8
- import io
9
- import base64
10
- from PIL import Image
11
- import gradio as gr
12
- import time
13
- import hashlib
14
- import requests
15
-
16
- from utils import build_logger
17
- from conversation import conv_seed_llama2
18
-
19
- IMG_FLAG = '<image>'
20
- LOGDIR = 'log'
21
-
22
- logger = build_logger("gradio_seed_x", LOGDIR)
23
- headers = {"User-Agent": "SEED-X Client"}
24
-
25
- no_change_btn = gr.Button.update()
26
- enable_btn = gr.Button.update(interactive=True)
27
- disable_btn = gr.Button.update(interactive=False)
28
-
29
-
30
- @dataclass
31
- class Arguments:
32
- server_port: Optional[int] = field(default=7860, metadata={"help": "network port"})
33
- server_name: Optional[str] = field(default='0.0.0.0', metadata={"help": "network address"})
34
- request_address: Optional[str] = field(default='http://127.0.0.1:7890/generate',
35
- metadata={"help": "request address"})
36
-
37
-
38
- parser = transformers.HfArgumentParser(Arguments)
39
- args, = parser.parse_args_into_dataclasses()
40
- conv_seed_llama = conv_seed_llama2
41
-
42
-
43
- def decode_image(encoded_image: str) -> Image:
44
- decoded_bytes = base64.b64decode(encoded_image.encode('utf-8'))
45
- buffer = io.BytesIO(decoded_bytes)
46
- image = Image.open(buffer)
47
- return image
48
-
49
-
50
- def encode_image(image: Image.Image, format: str = 'PNG') -> str:
51
- with io.BytesIO() as buffer:
52
- image.save(buffer, format=format)
53
- encoded_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
54
- return encoded_image
55
-
56
-
57
- def get_conv_log_filename():
58
- t = datetime.datetime.now()
59
- name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
60
- return name
61
-
62
-
63
- def get_conv_image_dir():
64
- name = os.path.join(LOGDIR, 'images')
65
- os.makedirs(name, exist_ok=True)
66
- return name
67
-
68
-
69
- def get_image_name(image, image_dir=None):
70
- buffer = io.BytesIO()
71
- image.save(buffer, format='PNG')
72
- image_bytes = buffer.getvalue()
73
- md5 = hashlib.md5(image_bytes).hexdigest()
74
-
75
- if image_dir is not None:
76
- image_name = os.path.join(image_dir, md5 + '.png')
77
- else:
78
- image_name = md5 + '.png'
79
-
80
- return image_name
81
-
82
-
83
- def resize_image_square(image, target_size=448):
84
- resized_image = image.resize((target_size, target_size))
85
- return resized_image
86
-
87
-
88
- def resize_image(image, max_size=512):
89
- width, height = image.size
90
- aspect_ratio = float(width) / float(height)
91
-
92
- if width > height:
93
- new_width = max_size
94
- new_height = int(new_width / aspect_ratio)
95
- else:
96
- new_height = max_size
97
- new_width = int(new_height * aspect_ratio)
98
-
99
- resized_image = image.resize((new_width, new_height))
100
- return resized_image
101
-
102
-
103
- def center_crop_image(image, max_aspect_ratio=1.5):
104
- width, height = image.size
105
- aspect_ratio = max(width, height) / min(width, height)
106
-
107
- if aspect_ratio >= max_aspect_ratio:
108
- if width > height:
109
- new_width = int(height * max_aspect_ratio)
110
- left = (width - new_width) // 2
111
- right = (width + new_width) // 2
112
- top = 0
113
- bottom = height
114
- else:
115
- new_height = int(width * max_aspect_ratio)
116
- left = 0
117
- right = width
118
- top = (height - new_height) // 2
119
- bottom = (height + new_height) // 2
120
-
121
- cropped_image = image.crop((left, top, right, bottom))
122
- return cropped_image
123
- else:
124
- return image
125
-
126
-
127
- def vote_last_response(state, vote_type, request: gr.Request):
128
- with open(get_conv_log_filename(), "a") as fout:
129
- data = {
130
- "tstamp": round(time.time(), 4),
131
- "type": vote_type,
132
- "state": state.dict(),
133
- "ip": request.client.host,
134
- }
135
- fout.write(json.dumps(data) + "\n")
136
-
137
-
138
- def upvote_last_response(state, request: gr.Request):
139
- logger.info(f"upvote. ip: {request.client.host}")
140
- vote_last_response(state, "upvote", request)
141
- return (disable_btn,) * 2
142
-
143
-
144
- def downvote_last_response(state, request: gr.Request):
145
- logger.info(f"downvote. ip: {request.client.host}")
146
- vote_last_response(state, "downvote", request)
147
- return (disable_btn,) * 2
148
-
149
-
150
- def regenerate(dialog_state, request: gr.Request):
151
- logger.info(f"regenerate. ip: {request.client.host}")
152
- if dialog_state.messages[-1]['role'] == dialog_state.roles[1]:
153
- dialog_state.messages.pop()
154
- return (
155
- dialog_state,
156
- dialog_state.to_gradio_chatbot(),
157
- ) + (disable_btn,) * 4
158
-
159
-
160
- def clear_history(request: gr.Request):
161
- logger.info(f"clear_history. ip: {request.client.host}")
162
- dialog_state = conv_seed_llama.copy()
163
- input_state = init_input_state()
164
- return (dialog_state, input_state, dialog_state.to_gradio_chatbot()) + (disable_btn,) * 4
165
-
166
-
167
- def init_input_state():
168
- return {'images': [], 'text': ''}
169
-
170
-
171
- def add_text(dialog_state, input_state, text, request: gr.Request):
172
- logger.info(f"add_text. ip: {request.client.host}.")
173
- # if len(input_state['text']) == 0:
174
- if text is None or len(text) == 0:
175
- # dialog_state.skip_next = True
176
- return (dialog_state, input_state, "", dialog_state.to_gradio_chatbot()) + (no_change_btn,) * 4
177
- input_state['text'] += text
178
-
179
-
180
- if len(dialog_state.messages) > 0 and dialog_state.messages[-1]['role'] == dialog_state.roles[0]:
181
- dialog_state.messages[-1]['message'] = input_state
182
- else:
183
- dialog_state.messages.append({'role': dialog_state.roles[0], 'message': input_state})
184
- print('add_text: ', dialog_state.to_gradio_chatbot())
185
-
186
- return (dialog_state, input_state, "", dialog_state.to_gradio_chatbot()) + (disable_btn,) * 4
187
-
188
-
189
- def is_blank(image):
190
- image_array = np.array(image)
191
- unique_colors = np.unique(image_array)
192
- print('unique_colors', len(unique_colors))
193
- return len(unique_colors) == 1
194
-
195
-
196
- def add_image(dialog_state, input_state, image, request: gr.Request):
197
- logger.info(f"add_image. ip: {request.client.host}.")
198
- if image is None:
199
- return (dialog_state, input_state, None, dialog_state.to_gradio_chatbot()) + (no_change_btn,) * 4
200
-
201
- image = image.convert('RGB')
202
-
203
- print('image size:', image.size)
204
-
205
- image = center_crop_image(image, max_aspect_ratio=10)
206
-
207
- image_dir = get_conv_image_dir()
208
- image_path = get_image_name(image=image, image_dir=image_dir)
209
- if not os.path.exists(image_path):
210
- image.save(image_path)
211
- input_state['images'].append(image_path)
212
- input_state['text'] += IMG_FLAG
213
-
214
- if len(dialog_state.messages) > 0 and dialog_state.messages[-1]['role'] == dialog_state.roles[0]:
215
- dialog_state.messages[-1]['message'] = input_state
216
- else:
217
- dialog_state.messages.append({'role': dialog_state.roles[0], 'message': input_state})
218
-
219
- print('add_image:', dialog_state)
220
-
221
- return (dialog_state, input_state, None, dialog_state.to_gradio_chatbot()) + (disable_btn,) * 4
222
-
223
-
224
- def http_bot(dialog_state, input_state, max_new_tokens, max_turns, force_image_gen, force_bbox,
225
- request: gr.Request):
226
- logger.info(f"http_bot. ip: {request.client.host}")
227
- print('input_state:', input_state)
228
-
229
- if len(dialog_state.messages) == 0 or dialog_state.messages[-1]['role'] != dialog_state.roles[0] or len(
230
- dialog_state.messages[-1]['message']['text'].strip(' ?.;!/')) == 0:
231
- return (dialog_state, input_state, dialog_state.to_gradio_chatbot()) + (no_change_btn,) * 4
232
-
233
- if len(dialog_state.messages) > max_turns * 2:
234
- output_state = init_input_state()
235
- output_state['text'] = 'Error: History exceeds maximum rounds, please clear history and restart.'
236
- dialog_state.messages.append({'role': dialog_state.roles[1], 'message': output_state})
237
- input_state = init_input_state()
238
- return (dialog_state, input_state, dialog_state.to_gradio_chatbot()) + (disable_btn,) * 3 + (enable_btn,)
239
-
240
- prompt = dialog_state.get_prompt()
241
- payload = {
242
- 'text': prompt['text'],
243
- 'max_new_tokens': int(max_new_tokens),
244
- 'images': prompt['images'],
245
- 'force_boi': force_image_gen,
246
- 'force_bbox': force_bbox,
247
- }
248
-
249
- print(
250
- 'request: ', {
251
- 'text': prompt['text'],
252
- 'max_new_tokens': int(max_new_tokens),
253
- })
254
- print('request_address', args.request_address)
255
- response = requests.request(method="POST", url=args.request_address, headers=headers, json=payload)
256
- results = response.json()
257
- print('response: ', {'text': results['text'], 'error_msg': results['error_msg']})
258
-
259
- output_state = init_input_state()
260
- image_dir = get_conv_image_dir()
261
- output_state['text'] = results['text']
262
-
263
- for image_base64 in results['images']:
264
- if image_base64 == '':
265
- image_path = ''
266
- else:
267
- image = decode_image(image_base64)
268
- image = image.convert('RGB')
269
- image_path = get_image_name(image=image, image_dir=image_dir)
270
- if not os.path.exists(image_path):
271
- image.save(image_path)
272
- output_state['images'].append(image_path)
273
-
274
- dialog_state.messages.append({'role': dialog_state.roles[1], 'message': output_state})
275
-
276
- vote_last_response(dialog_state, 'common', request)
277
- input_state = init_input_state()
278
- chatbot = update_error_msg(dialog_state.to_gradio_chatbot(), results['error_msg'])
279
- return (dialog_state, input_state, chatbot) + (enable_btn,) * 4
280
-
281
-
282
- def update_error_msg(chatbot, error_msg):
283
- if len(error_msg) > 0:
284
- info = '\n-------------\nSome errors occurred during response, please clear history and restart.\n' + '\n'.join(
285
- error_msg)
286
- chatbot[-1][-1] = chatbot[-1][-1] + info
287
-
288
- return chatbot
289
-
290
-
291
- def load_demo(request: gr.Request):
292
- logger.info(f"load_demo. ip: {request.client.host}")
293
- dialog_state = conv_seed_llama.copy()
294
- input_state = init_input_state()
295
- return dialog_state, input_state
296
-
297
-
298
- title = ("""
299
- # SEED-X-I
300
- [[Paper]](https://arxiv.org/abs/2404.14396) [[Code]](https://github.com/AILab-CVC/SEED-X)
301
-
302
- Demo of a general instruction-tuned model SEED-X-I (17B) from the foundation model SEED-X.
303
-
304
- SEED-X-I can follow multimodal instruction (including images with **dynamic resolutions**) and make responses with **images, texts and bounding boxes** in multi-turn conversation.
305
-
306
- SEED-X-I **does not support image manipulation**. If you want to experience **SEED-X-Edit** for high-precision image editing, please refer to [[Inference Code]](https://github.com/AILab-CVC/SEED-X).
307
-
308
- Due to insufficient GPU memory, when generating images, we need to offload the LLM to the CPU and move the de-tokenizer to the CPU, which will **result in a long processing time**. If you want to experience the normal model inference speed, you can run [[Inference Code]](https://github.com/AILab-CVC/SEED-X) locally.
309
-
310
-
311
- ## Tips:
312
- * Check out the conversation examples (at the bottom) for inspiration.
313
-
314
- * You can adjust "Max History Rounds" to try a conversation with up to five rounds. For more turns, you can download our checkpoints from GitHub and deploy them locally for inference.
315
-
316
- * Our demo supports a mix of images and texts as input. You can freely upload an image or enter text, and then click on "Add Image/Text". You can repeat the former step multiple times, and click on "Submit" for model inference at last.
317
-
318
- * You can click "Force Image Generation" to compel the model to produce images when necessary. For example, our model might struggle to generate images when there is an excessive amount of text-only context.
319
-
320
- * You can click "Force Bounding Box" to compel the model to produce bounding box for object detection.
321
-
322
- * SEED-X was trained with English-only data. It may process with other languages due to the inherent capabilities from LLaMA, but might not stable.
323
-
324
- """)
325
-
326
- css = """
327
- img {
328
- font-family: 'Helvetica';
329
- font-weight: 300;
330
- line-height: 2;
331
- text-align: center;
332
-
333
- width: auto;
334
- height: auto;
335
- display: block;
336
- position: relative;
337
- }
338
-
339
- img:before {
340
- content: " ";
341
- display: block;
342
-
343
- position: absolute;
344
- top: -10px;
345
- left: 0;
346
- height: calc(100% + 10px);
347
- width: 100%;
348
- background-color: rgb(230, 230, 230);
349
- border: 2px dotted rgb(200, 200, 200);
350
- border-radius: 5px;
351
- }
352
-
353
- img:after {
354
- content: " ";
355
- display: block;
356
- font-size: 16px;
357
- font-style: normal;
358
- font-family: FontAwesome;
359
- color: rgb(100, 100, 100);
360
-
361
- position: absolute;
362
- top: 5px;
363
- left: 0;
364
- width: 100%;
365
- text-align: center;
366
- }
367
-
368
- """
369
-
370
- if __name__ == '__main__':
371
-
372
- examples_mix = [
373
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/bank.png?raw=true', 'Can I conntect with an advisor on Sunday?'],
374
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/ground.png?raw=true',
375
- 'Is there anything in the image that can protect me from catching the flu virus when I go out? Show me the location.'],
376
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/arrow.jpg?raw=true', 'What is the object pointed by the red arrow?'],
377
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/shanghai.png?raw=true', 'Where was this image taken? Explain your answer.'],
378
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/GPT4.png?raw=true', 'How long does it take to make GPT-4 safer?'],
379
- ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/twitter.png?raw=true',
380
- 'Please provide a comprehensive description of this image.'],
381
- ]
382
-
383
- examples_text = [
384
- ['I want to build a two story cabin in the woods, with many commanding windows. Can you show me a picture?'],
385
- ['Use your imagination to design a concept image for Artificial General Intelligence (AGI). Show me an image.'],
386
- [
387
- 'Can you design an illustration for “The Three-Body Problem” to depict a scene from the novel? Show me a picture.'],
388
- [
389
- 'My four year old son loves toy trains. Can you design a fancy birthday cake for him? Please generate a picture.'],
390
- [
391
- 'Generate an image of a portrait of young nordic girl, age 25, freckled skin, neck tatoo, blue eyes 35mm lens, photography, ultra details.'],
392
- ['Generate an impressionist painting of an astronaut in a jungle.']
393
- ]
394
- with gr.Blocks(css=css) as demo:
395
- gr.Markdown(title)
396
- dialog_state = gr.State()
397
- input_state = gr.State()
398
- with gr.Row():
399
- with gr.Column(scale=3):
400
- with gr.Row():
401
- image = gr.Image(type='pil', label='input_image')
402
- with gr.Row():
403
- text = gr.Textbox(lines=5,
404
- show_label=False,
405
- label='input_text',
406
- elem_id='textbox',
407
- placeholder="Enter text or add image, and press submit,").style(container=False)
408
- with gr.Row():
409
- add_image_btn = gr.Button("Add Image")
410
- add_text_btn = gr.Button("Add Text")
411
-
412
- submit_btn = gr.Button("Submit")
413
-
414
- with gr.Row():
415
- max_new_tokens = gr.Slider(minimum=64,
416
- maximum=1024,
417
- value=768,
418
- step=64,
419
- interactive=True,
420
- label="Max Output Tokens")
421
- max_turns = gr.Slider(minimum=1, maximum=9, value=3, step=1, interactive=True,
422
- label="Max History Rounds")
423
- force_img_gen = gr.Radio(choices=[True, False], value=False, label='Force Image Generation')
424
- force_bbox = gr.Radio(choices=[True, False], value=False, label='Force Bounding Box')
425
-
426
- with gr.Column(scale=7):
427
- chatbot = gr.Chatbot(elem_id='chatbot', label="SEED-X-I").style(height=700)
428
- with gr.Row():
429
- upvote_btn = gr.Button(value="👍 Upvote", interactive=False)
430
- downvote_btn = gr.Button(value="👎 Downvote", interactive=False)
431
- regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False)
432
- clear_btn = gr.Button(value="🗑️ Clear history", interactive=False)
433
-
434
- with gr.Row():
435
- with gr.Column(scale=0.7):
436
- gr.Examples(examples=examples_mix, label='Input examples', inputs=[image, text])
437
- with gr.Column(scale=0.3):
438
- gr.Examples(examples=examples_text, label='Input examples', inputs=[text])
439
-
440
- # Register listeners
441
- btn_list = [upvote_btn, downvote_btn, regenerate_btn, clear_btn]
442
- upvote_btn.click(upvote_last_response, [dialog_state], [upvote_btn, downvote_btn])
443
- downvote_btn.click(downvote_last_response, [dialog_state], [upvote_btn, downvote_btn])
444
-
445
- regenerate_btn.click(regenerate, [dialog_state], [dialog_state, chatbot] + btn_list).then(
446
- http_bot, [dialog_state, input_state, max_new_tokens, max_turns, force_img_gen, force_bbox],
447
- [dialog_state, input_state, chatbot] + btn_list)
448
- add_image_btn.click(add_image, [dialog_state, input_state, image],
449
- [dialog_state, input_state, image, chatbot] + btn_list)
450
-
451
- add_text_btn.click(add_text, [dialog_state, input_state, text],
452
- [dialog_state, input_state, text, chatbot] + btn_list)
453
-
454
- submit_btn.click(
455
- add_image, [dialog_state, input_state, image], [dialog_state, input_state, image, chatbot] + btn_list).then(
456
- add_text, [dialog_state, input_state, text],
457
- [dialog_state, input_state, text, chatbot, upvote_btn, downvote_btn, regenerate_btn, clear_btn]).then(
458
- http_bot,
459
- [dialog_state, input_state, max_new_tokens, max_turns, force_img_gen, force_bbox],
460
- [dialog_state, input_state, chatbot] + btn_list)
461
- clear_btn.click(clear_history, None, [dialog_state, input_state, chatbot] + btn_list)
462
-
463
- demo.load(load_demo, None, [dialog_state, input_state])
464
-
465
- demo.launch(server_name=args.server_name, server_port=args.server_port, enable_queue=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/demo/utils.py DELETED
@@ -1,83 +0,0 @@
1
- import datetime
2
- import logging
3
- import logging.handlers
4
- import os
5
- import sys
6
-
7
- handler = None
8
-
9
-
10
- def build_logger(logger_name, logger_dir):
11
- global handler
12
-
13
- formatter = logging.Formatter(
14
- fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
15
- datefmt="%Y-%m-%d %H:%M:%S",
16
- )
17
-
18
- # Set the format of root handlers
19
- if not logging.getLogger().handlers:
20
- logging.basicConfig(level=logging.INFO)
21
- logging.getLogger().handlers[0].setFormatter(formatter)
22
-
23
- # Redirect stdout and stderr to loggers
24
- stdout_logger = logging.getLogger("stdout")
25
- stdout_logger.setLevel(logging.INFO)
26
- sl = StreamToLogger(stdout_logger, logging.INFO)
27
- sys.stdout = sl
28
-
29
- stderr_logger = logging.getLogger("stderr")
30
- stderr_logger.setLevel(logging.ERROR)
31
- sl = StreamToLogger(stderr_logger, logging.ERROR)
32
- sys.stderr = sl
33
-
34
- # Get logger
35
- logger = logging.getLogger(logger_name)
36
- logger.setLevel(logging.INFO)
37
-
38
- # Add a file handler for all loggers
39
- if handler is None:
40
- os.makedirs(logger_dir, exist_ok=True)
41
- filename = os.path.join(logger_dir, logger_name + '.log')
42
- handler = logging.handlers.TimedRotatingFileHandler(filename, when='D', utc=True)
43
- handler.setFormatter(formatter)
44
-
45
- for name, item in logging.root.manager.loggerDict.items():
46
- if isinstance(item, logging.Logger):
47
- item.addHandler(handler)
48
-
49
- return logger
50
-
51
-
52
- class StreamToLogger(object):
53
- """
54
- Fake file-like stream object that redirects writes to a logger instance.
55
- """
56
-
57
- def __init__(self, logger, log_level=logging.INFO):
58
- self.terminal = sys.stdout
59
- self.logger = logger
60
- self.log_level = log_level
61
- self.linebuf = ''
62
-
63
- def __getattr__(self, attr):
64
- return getattr(self.terminal, attr)
65
-
66
- def write(self, buf):
67
- temp_linebuf = self.linebuf + buf
68
- self.linebuf = ''
69
- for line in temp_linebuf.splitlines(True):
70
- # From the io.TextIOWrapper docs:
71
- # On output, if newline is None, any '\n' characters written
72
- # are translated to the system default line separator.
73
- # By default sys.stdout.write() expects '\n' newlines and then
74
- # translates them so this is still cross platform.
75
- if line[-1] == '\n':
76
- self.logger.log(self.log_level, line.rstrip())
77
- else:
78
- self.linebuf += line
79
-
80
- def flush(self):
81
- if self.linebuf != '':
82
- self.logger.log(self.log_level, self.linebuf.rstrip())
83
- self.linebuf = ''