Kate-03/GenAI-Arena
0
1import datetime2import time3import json4import uuid5import gradio as gr6import regex as re7from pathlib import Path8from .utils import *9from .log_utils import build_logger10from .constants import IMAGE_DIR, VIDEO_DIR11import imageio12from diffusers.utils import load_image13import torch14 15ig_logger = build_logger("gradio_web_server_image_generation", "gr_web_image_generation.log") # ig = image generation, loggers for single model direct chat16igm_logger = build_logger("gradio_web_server_image_generation_multi", "gr_web_image_generation_multi.log") # igm = image generation multi, loggers for side-by-side and battle17ie_logger = build_logger("gradio_web_server_image_editing", "gr_web_image_editing.log") # ie = image editing, loggers for single model direct chat18iem_logger = build_logger("gradio_web_server_image_editing_multi", "gr_web_image_editing_multi.log") # iem = image editing multi, loggers for side-by-side and battle19vg_logger = build_logger("gradio_web_server_video_generation", "gr_web_video_generation.log") # vg = video generation, loggers for single model direct chat20vgm_logger = build_logger("gradio_web_server_video_generation_multi", "gr_web_video_generation_multi.log") # vgm = video generation multi, loggers for side-by-side and battle21 22def save_any_image(image_file, file_path):23 if isinstance(image_file, str):24 image = load_image(image_file)25 image.save(file_path, 'JPEG')26 else:27 image_file.save(file_path, 'JPEG')28 29def vote_last_response_ig(state, vote_type, model_selector, request: gr.Request):30 with open(get_conv_log_filename(), "a") as fout:31 data = {32 "tstamp": round(time.time(), 4),33 "type": vote_type,34 "model": model_selector,35 "state": state.dict(),36 "ip": get_ip(request),37 }38 fout.write(json.dumps(data) + "\n")39 append_json_item_on_log_server(data, get_conv_log_filename())40 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'41 with open(output_file, 'w') as f:42 save_any_image(state.output, f)43 save_image_file_on_log_server(output_file)44 45def vote_last_response_igm(states, vote_type, model_selectors, request: gr.Request):46 with open(get_conv_log_filename(), "a") as fout:47 data = {48 "tstamp": round(time.time(), 4),49 "type": vote_type,50 "models": [x for x in model_selectors],51 "states": [x.dict() for x in states],52 "ip": get_ip(request),53 }54 fout.write(json.dumps(data) + "\n")55 append_json_item_on_log_server(data, get_conv_log_filename())56 for state in states:57 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'58 with open(output_file, 'w') as f:59 save_any_image(state.output, f)60 save_image_file_on_log_server(output_file)61 62def vote_last_response_ie(state, vote_type, model_selector, request: gr.Request):63 with open(get_conv_log_filename(), "a") as fout:64 data = {65 "tstamp": round(time.time(), 4),66 "type": vote_type,67 "model": model_selector,68 "state": state.dict(),69 "ip": get_ip(request),70 }71 fout.write(json.dumps(data) + "\n")72 append_json_item_on_log_server(data, get_conv_log_filename())73 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}.jpg'74 source_file = f'{IMAGE_DIR}/edition/{state.conv_id}_source.jpg'75 with open(output_file, 'w') as f:76 save_any_image(state.output, f)77 with open(source_file, 'w') as sf:78 save_any_image(state.source_image, sf)79 save_image_file_on_log_server(output_file)80 save_image_file_on_log_server(source_file)81 82def vote_last_response_iem(states, vote_type, model_selectors, request: gr.Request):83 with open(get_conv_log_filename(), "a") as fout:84 data = {85 "tstamp": round(time.time(), 4),86 "type": vote_type,87 "models": [x for x in model_selectors],88 "states": [x.dict() for x in states],89 "ip": get_ip(request),90 }91 fout.write(json.dumps(data) + "\n")92 append_json_item_on_log_server(data, get_conv_log_filename())93 for state in states:94 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}.jpg'95 source_file = f'{IMAGE_DIR}/edition/{state.conv_id}_source.jpg'96 with open(output_file, 'w') as f:97 save_any_image(state.output, f)98 with open(source_file, 'w') as sf:99 save_any_image(state.source_image, sf)100 save_image_file_on_log_server(output_file)101 save_image_file_on_log_server(source_file)102 103 104def vote_last_response_vg(state, vote_type, model_selector, request: gr.Request):105 with open(get_conv_log_filename(), "a") as fout:106 data = {107 "tstamp": round(time.time(), 4),108 "type": vote_type,109 "model": model_selector,110 "state": state.dict(),111 "ip": get_ip(request),112 }113 fout.write(json.dumps(data) + "\n")114 append_json_item_on_log_server(data, get_conv_log_filename())115 116 output_file = f'{VIDEO_DIR}/generation/{state.conv_id}.mp4'117 os.makedirs(os.path.dirname(output_file), exist_ok=True)118 if state.model_name.startswith('fal'):119 r = requests.get(state.output)120 with open(output_file, 'wb') as outfile:121 outfile.write(r.content)122 else:123 print("======== video shape: ========")124 print(state.output.shape)125 # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]126 if state.output.shape[-1] != 3:127 state.output = state.output.permute(0, 2, 3, 1)128 imageio.mimwrite(output_file, state.output, fps=8, quality=9)129 save_video_file_on_log_server(output_file)130 131 132 133def vote_last_response_vgm(states, vote_type, model_selectors, request: gr.Request):134 with open(get_conv_log_filename(), "a") as fout:135 data = {136 "tstamp": round(time.time(), 4),137 "type": vote_type,138 "models": [x for x in model_selectors],139 "states": [x.dict() for x in states],140 "ip": get_ip(request),141 }142 fout.write(json.dumps(data) + "\n")143 append_json_item_on_log_server(data, get_conv_log_filename())144 for state in states:145 output_file = f'{VIDEO_DIR}/generation/{state.conv_id}.mp4'146 os.makedirs(os.path.dirname(output_file), exist_ok=True)147 if state.model_name.startswith('fal'):148 r = requests.get(state.output)149 with open(output_file, 'wb') as outfile:150 outfile.write(r.content)151 elif isinstance(state.output, torch.Tensor):152 print("======== video shape: ========")153 print(state.output.shape)154 # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]155 if state.output.shape[-1] != 3:156 state.output = state.output.permute(0, 2, 3, 1)157 imageio.mimwrite(output_file, state.output, fps=8, quality=9)158 else:159 r = requests.get(state.output)160 with open(output_file, 'wb') as outfile:161 outfile.write(r.content)162 save_video_file_on_log_server(output_file)163 164 165## Image Generation (IG) Single Model Direct Chat166def upvote_last_response_ig(state, model_selector, request: gr.Request):167 ip = get_ip(request)168 ig_logger.info(f"upvote. ip: {ip}")169 vote_last_response_ig(state, "upvote", model_selector, request)170 return ("",) + (disable_btn,) * 3171 172def downvote_last_response_ig(state, model_selector, request: gr.Request):173 ip = get_ip(request)174 ig_logger.info(f"downvote. ip: {ip}")175 vote_last_response_ig(state, "downvote", model_selector, request)176 return ("",) + (disable_btn,) * 3177 178 179def flag_last_response_ig(state, model_selector, request: gr.Request):180 ip = get_ip(request)181 ig_logger.info(f"flag. ip: {ip}")182 vote_last_response_ig(state, "flag", model_selector, request)183 return ("",) + (disable_btn,) * 3184 185## Image Generation Multi (IGM) Side-by-Side and Battle186 187def leftvote_last_response_igm(188 state0, state1, model_selector0, model_selector1, request: gr.Request189):190 igm_logger.info(f"leftvote (named). ip: {get_ip(request)}")191 vote_last_response_igm(192 [state0, state1], "leftvote", [model_selector0, model_selector1], request193 )194 if model_selector0 == "":195 return ("",) + (disable_btn,) * 4 + (196 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),197 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))198 else:199 return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),200 gr.Markdown(state1.model_name, visible=True))201 202def rightvote_last_response_igm(203 state0, state1, model_selector0, model_selector1, request: gr.Request204):205 igm_logger.info(f"rightvote (named). ip: {get_ip(request)}")206 vote_last_response_igm(207 [state0, state1], "rightvote", [model_selector0, model_selector1], request208 )209 print(model_selector0)210 if model_selector0 == "":211 return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))212 else:213 return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),214 gr.Markdown(state1.model_name, visible=True))215 216 217def tievote_last_response_igm(218 state0, state1, model_selector0, model_selector1, request: gr.Request219):220 igm_logger.info(f"tievote (named). ip: {get_ip(request)}")221 vote_last_response_igm(222 [state0, state1], "tievote", [model_selector0, model_selector1], request223 )224 if model_selector0 == "":225 return ("",) + (disable_btn,) * 4 + (226 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),227 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))228 else:229 return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),230 gr.Markdown(state1.model_name, visible=True))231 232 233def bothbad_vote_last_response_igm(234 state0, state1, model_selector0, model_selector1, request: gr.Request235):236 igm_logger.info(f"bothbad_vote (named). ip: {get_ip(request)}")237 vote_last_response_igm(238 [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request239 )240 if model_selector0 == "":241 return ("",) + (disable_btn,) * 4 + (242 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),243 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))244 else:245 return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),246 gr.Markdown(state1.model_name, visible=True))247 248## Image Editing (IE) Single Model Direct Chat249 250def upvote_last_response_ie(state, model_selector, request: gr.Request):251 ip = get_ip(request)252 ie_logger.info(f"upvote. ip: {ip}")253 vote_last_response_ie(state, "upvote", model_selector, request)254 return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3255 256def downvote_last_response_ie(state, model_selector, request: gr.Request):257 ip = get_ip(request)258 ie_logger.info(f"downvote. ip: {ip}")259 vote_last_response_ie(state, "downvote", model_selector, request)260 return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3261 262def flag_last_response_ie(state, model_selector, request: gr.Request):263 ip = get_ip(request)264 ie_logger.info(f"flag. ip: {ip}")265 vote_last_response_ie(state, "flag", model_selector, request)266 return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3267 268## Image Editing Multi (IEM) Side-by-Side and Battle269def leftvote_last_response_iem(270 state0, state1, model_selector0, model_selector1, request: gr.Request271):272 iem_logger.info(f"leftvote (anony). ip: {get_ip(request)}")273 vote_last_response_iem(274 [state0, state1], "leftvote", [model_selector0, model_selector1], request275 )276 # names = (277 # "### Model A: " + state0.model_name,278 # "### Model B: " + state1.model_name,279 # )280 # names = (state0.model_name, state1.model_name)281 if model_selector0 == "":282 names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))283 else:284 names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))285 return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4286 287def rightvote_last_response_iem(288 state0, state1, model_selector0, model_selector1, request: gr.Request289):290 iem_logger.info(f"rightvote (anony). ip: {get_ip(request)}")291 vote_last_response_iem(292 [state0, state1], "rightvote", [model_selector0, model_selector1], request293 )294 # names = (295 # "### Model A: " + state0.model_name,296 # "### Model B: " + state1.model_name,297 # )298 if model_selector0 == "":299 names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),300 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))301 else:302 names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))303 return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4304 305def tievote_last_response_iem(306 state0, state1, model_selector0, model_selector1, request: gr.Request307):308 iem_logger.info(f"tievote (anony). ip: {get_ip(request)}")309 vote_last_response_iem(310 [state0, state1], "tievote", [model_selector0, model_selector1], request311 )312 if model_selector0 == "":313 names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),314 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))315 else:316 names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))317 return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4318 319def bothbad_vote_last_response_iem(320 state0, state1, model_selector0, model_selector1, request: gr.Request321):322 iem_logger.info(f"bothbad_vote (anony). ip: {get_ip(request)}")323 vote_last_response_iem(324 [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request325 )326 if model_selector0 == "":327 names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),328 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))329 else:330 names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))331 return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4332 333 334## Video Generation (VG) Single Model Direct Chat335def upvote_last_response_vg(state, model_selector, request: gr.Request):336 ip = get_ip(request)337 vg_logger.info(f"upvote. ip: {ip}")338 vote_last_response_vg(state, "upvote", model_selector, request)339 return ("",) + (disable_btn,) * 3340 341def downvote_last_response_vg(state, model_selector, request: gr.Request):342 ip = get_ip(request)343 vg_logger.info(f"downvote. ip: {ip}")344 vote_last_response_vg(state, "downvote", model_selector, request)345 return ("",) + (disable_btn,) * 3346 347 348def flag_last_response_vg(state, model_selector, request: gr.Request):349 ip = get_ip(request)350 vg_logger.info(f"flag. ip: {ip}")351 vote_last_response_vg(state, "flag", model_selector, request)352 return ("",) + (disable_btn,) * 3353 354## Image Generation Multi (IGM) Side-by-Side and Battle355 356def leftvote_last_response_vgm(357 state0, state1, model_selector0, model_selector1, request: gr.Request358):359 vgm_logger.info(f"leftvote (named). ip: {get_ip(request)}")360 vote_last_response_vgm(361 [state0, state1], "leftvote", [model_selector0, model_selector1], request362 )363 if model_selector0 == "":364 return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))365 else:366 return ("",) + (disable_btn,) * 4 + (367 gr.Markdown(state0.model_name, visible=False),368 gr.Markdown(state1.model_name, visible=False))369 370 371def rightvote_last_response_vgm(372 state0, state1, model_selector0, model_selector1, request: gr.Request373):374 vgm_logger.info(f"rightvote (named). ip: {get_ip(request)}")375 vote_last_response_vgm(376 [state0, state1], "rightvote", [model_selector0, model_selector1], request377 )378 if model_selector0 == "":379 return ("",) + (disable_btn,) * 4 + (380 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),381 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))382 else:383 return ("",) + (disable_btn,) * 4 + (384 gr.Markdown(state0.model_name, visible=False),385 gr.Markdown(state1.model_name, visible=False))386 387def tievote_last_response_vgm(388 state0, state1, model_selector0, model_selector1, request: gr.Request389):390 vgm_logger.info(f"tievote (named). ip: {get_ip(request)}")391 vote_last_response_vgm(392 [state0, state1], "tievote", [model_selector0, model_selector1], request393 )394 if model_selector0 == "":395 return ("",) + (disable_btn,) * 4 + (396 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),397 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))398 else:399 return ("",) + (disable_btn,) * 4 + (400 gr.Markdown(state0.model_name, visible=False),401 gr.Markdown(state1.model_name, visible=False))402 403 404def bothbad_vote_last_response_vgm(405 state0, state1, model_selector0, model_selector1, request: gr.Request406):407 vgm_logger.info(f"bothbad_vote (named). ip: {get_ip(request)}")408 vote_last_response_vgm(409 [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request410 )411 if model_selector0 == "":412 return ("",) + (disable_btn,) * 4 + (413 gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),414 gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))415 else:416 return ("",) + (disable_btn,) * 4 + (417 gr.Markdown(state0.model_name, visible=False),418 gr.Markdown(state1.model_name, visible=False))419 420share_js = """421function (a, b, c, d) {422 const captureElement = document.querySelector('#share-region-named');423 html2canvas(captureElement)424 .then(canvas => {425 canvas.style.display = 'none'426 document.body.appendChild(canvas)427 return canvas428 })429 .then(canvas => {430 const image = canvas.toDataURL('image/png')431 const a = document.createElement('a')432 a.setAttribute('download', 'chatbot-arena.png')433 a.setAttribute('href', image)434 a.click()435 canvas.remove()436 });437 return [a, b, c, d];438}439"""440def share_click_igm(state0, state1, model_selector0, model_selector1, request: gr.Request):441 igm_logger.info(f"share (anony). ip: {get_ip(request)}")442 if state0 is not None and state1 is not None:443 vote_last_response_igm(444 [state0, state1], "share", [model_selector0, model_selector1], request445 )446 447def share_click_iem(state0, state1, model_selector0, model_selector1, request: gr.Request):448 iem_logger.info(f"share (anony). ip: {get_ip(request)}")449 if state0 is not None and state1 is not None:450 vote_last_response_iem(451 [state0, state1], "share", [model_selector0, model_selector1], request452 )453 454## All Generation Gradio Interface455 456class ImageStateIG:457 def __init__(self, model_name):458 self.conv_id = uuid.uuid4().hex459 self.model_name = model_name460 self.prompt = None461 self.output = None462 463 def dict(self):464 base = {465 "conv_id": self.conv_id,466 "model_name": self.model_name,467 "prompt": self.prompt468 }469 return base470 471class ImageStateIE:472 def __init__(self, model_name):473 self.conv_id = uuid.uuid4().hex474 self.model_name = model_name475 self.source_prompt = None476 self.target_prompt = None477 self.instruct_prompt = None478 self.source_image = None479 self.output = None480 481 def dict(self):482 base = {483 "conv_id": self.conv_id,484 "model_name": self.model_name,485 "source_prompt": self.source_prompt,486 "target_prompt": self.target_prompt,487 "instruct_prompt": self.instruct_prompt488 }489 return base490 491class VideoStateVG:492 def __init__(self, model_name):493 self.conv_id = uuid.uuid4().hex494 self.model_name = model_name495 self.prompt = None496 self.output = None497 498 def dict(self):499 base = {500 "conv_id": self.conv_id,501 "model_name": self.model_name,502 "prompt": self.prompt503 }504 return base505 506 507def generate_ig(gen_func, state, text, model_name, request: gr.Request):508 if not text:509 raise gr.Warning("Prompt cannot be empty.")510 if not model_name:511 raise gr.Warning("Model name cannot be empty.")512 state = ImageStateIG(model_name)513 ip = get_ip(request)514 ig_logger.info(f"generate. ip: {ip}")515 start_tstamp = time.time()516 generated_image = gen_func(text, model_name)517 state.prompt = text518 state.output = generated_image519 state.model_name = model_name520 if generated_image == '':521 with open(get_nsfw_conv_log_filename(), "a") as fout:522 data = {523 "type": "chat",524 "model": model_name,525 "gen_params": {},526 "start": round(start_tstamp, 4),527 "state": state.dict(),528 "ip": get_ip(request),529 }530 fout.write(json.dumps(data) + "\n")531 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())532 raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")533 534 yield state, generated_image535 536 finish_tstamp = time.time()537 # logger.info(f"===output===: {output}")538 539 with open(get_conv_log_filename(), "a") as fout:540 data = {541 "tstamp": round(finish_tstamp, 4),542 "type": "chat",543 "model": model_name,544 "gen_params": {},545 "start": round(start_tstamp, 4),546 "finish": round(finish_tstamp, 4),547 "state": state.dict(),548 "ip": get_ip(request),549 }550 fout.write(json.dumps(data) + "\n")551 append_json_item_on_log_server(data, get_conv_log_filename())552 553 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'554 os.makedirs(os.path.dirname(output_file), exist_ok=True)555 with open(output_file, 'w') as f:556 save_any_image(state.output, f)557 save_image_file_on_log_server(output_file)558 559def generate_ig_museum(gen_func, state, model_name, request: gr.Request):560 if not model_name:561 raise gr.Warning("Model name cannot be empty.")562 state = ImageStateIG(model_name)563 ip = get_ip(request)564 ig_logger.info(f"generate. ip: {ip}")565 start_tstamp = time.time()566 generated_image, text = gen_func(model_name)567 state.prompt = text568 state.output = generated_image569 state.model_name = model_name570 571 yield state, generated_image, text572 573 finish_tstamp = time.time()574 # logger.info(f"===output===: {output}")575 576 with open(get_conv_log_filename(), "a") as fout:577 data = {578 "tstamp": round(finish_tstamp, 4),579 "type": "chat",580 "model": model_name,581 "gen_params": {},582 "start": round(start_tstamp, 4),583 "finish": round(finish_tstamp, 4),584 "state": state.dict(),585 "ip": get_ip(request),586 }587 fout.write(json.dumps(data) + "\n")588 append_json_item_on_log_server(data, get_conv_log_filename())589 590 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'591 os.makedirs(os.path.dirname(output_file), exist_ok=True)592 with open(output_file, 'w') as f:593 save_any_image(state.output, f)594 save_image_file_on_log_server(output_file)595 596def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):597 if not text:598 raise gr.Warning("Prompt cannot be empty.")599 if not model_name0:600 raise gr.Warning("Model name A cannot be empty.")601 if not model_name1:602 raise gr.Warning("Model name B cannot be empty.")603 state0 = ImageStateIG(model_name0)604 state1 = ImageStateIG(model_name1)605 ip = get_ip(request)606 igm_logger.info(f"generate. ip: {ip}")607 start_tstamp = time.time()608 # Remove ### Model (A|B): from model name609 model_name0 = re.sub(r"### Model A: ", "", model_name0)610 model_name1 = re.sub(r"### Model B: ", "", model_name1)611 generated_image0, generated_image1 = gen_func(text, model_name0, model_name1)612 state0.prompt = text613 state1.prompt = text614 state0.output = generated_image0615 state1.output = generated_image1616 state0.model_name = model_name0617 state1.model_name = model_name1618 if generated_image0 == '' and generated_image1 == '':619 with open(get_nsfw_conv_log_filename(), "a") as fout:620 data = {621 "type": "chat",622 "model": model_name0,623 "gen_params": {},624 "start": round(start_tstamp, 4),625 "state": state0.dict(),626 "ip": get_ip(request),627 }628 fout.write(json.dumps(data) + "\n")629 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())630 data = {631 "type": "chat",632 "model": model_name1,633 "gen_params": {},634 "start": round(start_tstamp, 4),635 "state": state1.dict(),636 "ip": get_ip(request),637 }638 fout.write(json.dumps(data) + "\n")639 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())640 raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")641 642 yield state0, state1, generated_image0, generated_image1643 644 finish_tstamp = time.time()645 # logger.info(f"===output===: {output}")646 647 with open(get_conv_log_filename(), "a") as fout:648 data = {649 "tstamp": round(finish_tstamp, 4),650 "type": "chat",651 "model": model_name0,652 "gen_params": {},653 "start": round(start_tstamp, 4),654 "finish": round(finish_tstamp, 4),655 "state": state0.dict(),656 "ip": get_ip(request),657 }658 fout.write(json.dumps(data) + "\n")659 append_json_item_on_log_server(data, get_conv_log_filename())660 data = {661 "tstamp": round(finish_tstamp, 4),662 "type": "chat",663 "model": model_name1,664 "gen_params": {},665 "start": round(start_tstamp, 4),666 "finish": round(finish_tstamp, 4),667 "state": state1.dict(),668 "ip": get_ip(request),669 }670 fout.write(json.dumps(data) + "\n")671 append_json_item_on_log_server(data, get_conv_log_filename())672 673 for i, state in enumerate([state0, state1]):674 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'675 os.makedirs(os.path.dirname(output_file), exist_ok=True)676 with open(output_file, 'w') as f:677 save_any_image(state.output, f)678 save_image_file_on_log_server(output_file)679 680def generate_igm_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):681 if not model_name0:682 raise gr.Warning("Model name A cannot be empty.")683 if not model_name1:684 raise gr.Warning("Model name B cannot be empty.")685 state0 = ImageStateIG(model_name0)686 state1 = ImageStateIG(model_name1)687 ip = get_ip(request)688 igm_logger.info(f"generate. ip: {ip}")689 start_tstamp = time.time()690 # Remove ### Model (A|B): from model name691 model_name0 = re.sub(r"### Model A: ", "", model_name0)692 model_name1 = re.sub(r"### Model B: ", "", model_name1)693 generated_image0, generated_image1, text = gen_func(model_name0, model_name1)694 state0.prompt = text695 state1.prompt = text696 state0.output = generated_image0697 state1.output = generated_image1698 state0.model_name = model_name0699 state1.model_name = model_name1700 701 yield state0, state1, generated_image0, generated_image1, text702 703 finish_tstamp = time.time()704 # logger.info(f"===output===: {output}")705 706 with open(get_conv_log_filename(), "a") as fout:707 data = {708 "tstamp": round(finish_tstamp, 4),709 "type": "chat",710 "model": model_name0,711 "gen_params": {},712 "start": round(start_tstamp, 4),713 "finish": round(finish_tstamp, 4),714 "state": state0.dict(),715 "ip": get_ip(request),716 }717 fout.write(json.dumps(data) + "\n")718 append_json_item_on_log_server(data, get_conv_log_filename())719 data = {720 "tstamp": round(finish_tstamp, 4),721 "type": "chat",722 "model": model_name1,723 "gen_params": {},724 "start": round(start_tstamp, 4),725 "finish": round(finish_tstamp, 4),726 "state": state1.dict(),727 "ip": get_ip(request),728 }729 fout.write(json.dumps(data) + "\n")730 append_json_item_on_log_server(data, get_conv_log_filename())731 732 for i, state in enumerate([state0, state1]):733 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'734 os.makedirs(os.path.dirname(output_file), exist_ok=True)735 with open(output_file, 'w') as f:736 save_any_image(state.output, f)737 save_image_file_on_log_server(output_file)738 739 740def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):741 if not text:742 raise gr.Warning("Prompt cannot be empty.")743 state0 = ImageStateIG(model_name0)744 state1 = ImageStateIG(model_name1)745 ip = get_ip(request)746 igm_logger.info(f"generate. ip: {ip}")747 start_tstamp = time.time()748 model_name0 = ""749 model_name1 = ""750 generated_image0, generated_image1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)751 state0.prompt = text752 state1.prompt = text753 state0.output = generated_image0754 state1.output = generated_image1755 state0.model_name = model_name0756 state1.model_name = model_name1757 if generated_image0 == '' and generated_image1 == '':758 with open(get_nsfw_conv_log_filename(), "a") as fout:759 data = {760 "type": "chat",761 "model": model_name0,762 "gen_params": {},763 "start": round(start_tstamp, 4),764 "state": state0.dict(),765 "ip": get_ip(request),766 }767 fout.write(json.dumps(data) + "\n")768 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())769 data = {770 "type": "chat",771 "model": model_name1,772 "gen_params": {},773 "start": round(start_tstamp, 4),774 "state": state1.dict(),775 "ip": get_ip(request),776 }777 fout.write(json.dumps(data) + "\n")778 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())779 raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")780 781 782 yield state0, state1, generated_image0, generated_image1, \783 gr.Markdown(f"### Model A: {model_name0}", visible=False), gr.Markdown(f"### Model B: {model_name1}", visible=False)784 785 finish_tstamp = time.time()786 # logger.info(f"===output===: {output}")787 788 with open(get_conv_log_filename(), "a") as fout:789 data = {790 "tstamp": round(finish_tstamp, 4),791 "type": "chat",792 "model": model_name0,793 "gen_params": {},794 "start": round(start_tstamp, 4),795 "finish": round(finish_tstamp, 4),796 "state": state0.dict(),797 "ip": get_ip(request),798 }799 fout.write(json.dumps(data) + "\n")800 append_json_item_on_log_server(data, get_conv_log_filename())801 data = {802 "tstamp": round(finish_tstamp, 4),803 "type": "chat",804 "model": model_name1,805 "gen_params": {},806 "start": round(start_tstamp, 4),807 "finish": round(finish_tstamp, 4),808 "state": state1.dict(),809 "ip": get_ip(request),810 }811 fout.write(json.dumps(data) + "\n")812 append_json_item_on_log_server(data, get_conv_log_filename())813 814 for i, state in enumerate([state0, state1]):815 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'816 os.makedirs(os.path.dirname(output_file), exist_ok=True)817 with open(output_file, 'w') as f:818 save_any_image(state.output, f)819 save_image_file_on_log_server(output_file)820 821def generate_igm_annoy_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):822 state0 = ImageStateIG(model_name0)823 state1 = ImageStateIG(model_name1)824 ip = get_ip(request)825 igm_logger.info(f"generate. ip: {ip}")826 start_tstamp = time.time()827 # model_name0 = re.sub(r"### Model A: ", "", model_name0)828 # model_name1 = re.sub(r"### Model B: ", "", model_name1)829 model_name0 = ""830 model_name1 = ""831 generated_image0, generated_image1, model_name0, model_name1, text = gen_func(model_name0, model_name1)832 state0.prompt = text833 state1.prompt = text834 state0.output = generated_image0835 state1.output = generated_image1836 state0.model_name = model_name0837 state1.model_name = model_name1838 839 yield state0, state1, generated_image0, generated_image1, text,\840 gr.Markdown(f"### Model A: {model_name0}", visible=False), gr.Markdown(f"### Model B: {model_name1}", visible=False)841 842 finish_tstamp = time.time()843 # logger.info(f"===output===: {output}")844 845 with open(get_conv_log_filename(), "a") as fout:846 data = {847 "tstamp": round(finish_tstamp, 4),848 "type": "chat",849 "model": model_name0,850 "gen_params": {},851 "start": round(start_tstamp, 4),852 "finish": round(finish_tstamp, 4),853 "state": state0.dict(),854 "ip": get_ip(request),855 }856 fout.write(json.dumps(data) + "\n")857 append_json_item_on_log_server(data, get_conv_log_filename())858 data = {859 "tstamp": round(finish_tstamp, 4),860 "type": "chat",861 "model": model_name1,862 "gen_params": {},863 "start": round(start_tstamp, 4),864 "finish": round(finish_tstamp, 4),865 "state": state1.dict(),866 "ip": get_ip(request),867 }868 fout.write(json.dumps(data) + "\n")869 append_json_item_on_log_server(data, get_conv_log_filename())870 871 for i, state in enumerate([state0, state1]):872 output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'873 os.makedirs(os.path.dirname(output_file), exist_ok=True)874 with open(output_file, 'w') as f:875 save_any_image(state.output, f)876 save_image_file_on_log_server(output_file)877 878def generate_ie(gen_func, state, source_text, target_text, instruct_text, source_image, model_name, request: gr.Request):879 if not source_text:880 raise gr.Warning("Source prompt cannot be empty.")881 if not target_text:882 raise gr.Warning("Target prompt cannot be empty.")883 if not instruct_text:884 raise gr.Warning("Instruction prompt cannot be empty.")885 if not source_image:886 raise gr.Warning("Source image cannot be empty.")887 if not model_name:888 raise gr.Warning("Model name cannot be empty.")889 state = ImageStateIE(model_name)890 ip = get_ip(request)891 ig_logger.info(f"generate. ip: {ip}")892 start_tstamp = time.time()893 generated_image = gen_func(source_text, target_text, instruct_text, source_image, model_name)894 state.source_prompt = source_text895 state.target_prompt = target_text896 state.instruct_prompt = instruct_text897 state.source_image = source_image898 state.output = generated_image899 state.model_name = model_name900 901 if generated_image == '':902 with open(get_nsfw_conv_log_filename(), "a") as fout:903 data = {904 "type": "chat",905 "model": model_name,906 "gen_params": {},907 "start": round(start_tstamp, 4),908 "state": state.dict(),909 "ip": get_ip(request),910 }911 fout.write(json.dumps(data) + "\n")912 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())913 raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")914 915 yield state, generated_image916 917 finish_tstamp = time.time()918 # logger.info(f"===output===: {output}")919 920 with open(get_conv_log_filename(), "a") as fout:921 data = {922 "tstamp": round(finish_tstamp, 4),923 "type": "chat",924 "model": model_name,925 "gen_params": {},926 "start": round(start_tstamp, 4),927 "finish": round(finish_tstamp, 4),928 "state": state.dict(),929 "ip": get_ip(request),930 }931 fout.write(json.dumps(data) + "\n")932 append_json_item_on_log_server(data, get_conv_log_filename())933 934 src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'935 os.makedirs(os.path.dirname(src_img_file), exist_ok=True)936 with open(src_img_file, 'w') as f:937 save_any_image(state.source_image, f)938 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'939 with open(output_file, 'w') as f:940 save_any_image(state.output, f)941 save_image_file_on_log_server(src_img_file)942 save_image_file_on_log_server(output_file)943 944def generate_ie_museum(gen_func, state, model_name, request: gr.Request):945 if not model_name:946 raise gr.Warning("Model name cannot be empty.")947 state = ImageStateIE(model_name)948 ip = get_ip(request)949 ig_logger.info(f"generate. ip: {ip}")950 start_tstamp = time.time()951 source_image, generated_image, source_text, target_text, instruct_text = gen_func(model_name)952 state.source_prompt = source_text953 state.target_prompt = target_text954 state.instruct_prompt = instruct_text955 state.source_image = source_image956 state.output = generated_image957 state.model_name = model_name958 959 yield state, generated_image, source_image, source_text, target_text, instruct_text960 961 finish_tstamp = time.time()962 # logger.info(f"===output===: {output}")963 964 with open(get_conv_log_filename(), "a") as fout:965 data = {966 "tstamp": round(finish_tstamp, 4),967 "type": "chat",968 "model": model_name,969 "gen_params": {},970 "start": round(start_tstamp, 4),971 "finish": round(finish_tstamp, 4),972 "state": state.dict(),973 "ip": get_ip(request),974 }975 fout.write(json.dumps(data) + "\n")976 append_json_item_on_log_server(data, get_conv_log_filename())977 978 src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'979 os.makedirs(os.path.dirname(src_img_file), exist_ok=True)980 with open(src_img_file, 'w') as f:981 save_any_image(state.source_image, f)982 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'983 with open(output_file, 'w') as f:984 save_any_image(state.output, f)985 save_image_file_on_log_server(src_img_file)986 save_image_file_on_log_server(output_file)987 988 989def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):990 if not source_text:991 raise gr.Warning("Source prompt cannot be empty.")992 if not target_text:993 raise gr.Warning("Target prompt cannot be empty.")994 if not instruct_text:995 raise gr.Warning("Instruction prompt cannot be empty.")996 if not source_image:997 raise gr.Warning("Source image cannot be empty.")998 if not model_name0:999 raise gr.Warning("Model name A cannot be empty.")1000 if not model_name1:1001 raise gr.Warning("Model name B cannot be empty.")1002 state0 = ImageStateIE(model_name0)1003 state1 = ImageStateIE(model_name1)1004 ip = get_ip(request)1005 igm_logger.info(f"generate. ip: {ip}")1006 start_tstamp = time.time()1007 model_name0 = re.sub(r"### Model A: ", "", model_name0)1008 model_name1 = re.sub(r"### Model B: ", "", model_name1)1009 generated_image0, generated_image1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)1010 state0.source_prompt = source_text1011 state0.target_prompt = target_text1012 state0.instruct_prompt = instruct_text1013 state0.source_image = source_image1014 state0.output = generated_image01015 state0.model_name = model_name01016 state1.source_prompt = source_text1017 state1.target_prompt = target_text1018 state1.instruct_prompt = instruct_text1019 state1.source_image = source_image1020 state1.output = generated_image11021 state1.model_name = model_name11022 1023 if generated_image0 == '' and generated_image1 == '':1024 with open(get_nsfw_conv_log_filename(), "a") as fout:1025 data = {1026 "type": "chat",1027 "model": model_name0,1028 "gen_params": {},1029 "start": round(start_tstamp, 4),1030 "state": state0.dict(),1031 "ip": get_ip(request),1032 }1033 fout.write(json.dumps(data) + "\n")1034 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())1035 data = {1036 "type": "chat",1037 "model": model_name1,1038 "gen_params": {},1039 "start": round(start_tstamp, 4),1040 "state": state1.dict(),1041 "ip": get_ip(request),1042 }1043 fout.write(json.dumps(data) + "\n")1044 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())1045 raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")1046 1047 1048 yield state0, state1, generated_image0, generated_image11049 1050 finish_tstamp = time.time()1051 # logger.info(f"===output===: {output}")1052 1053 with open(get_conv_log_filename(), "a") as fout:1054 data = {1055 "tstamp": round(finish_tstamp, 4),1056 "type": "chat",1057 "model": model_name0,1058 "gen_params": {},1059 "start": round(start_tstamp, 4),1060 "finish": round(finish_tstamp, 4),1061 "state": state0.dict(),1062 "ip": get_ip(request),1063 }1064 fout.write(json.dumps(data) + "\n")1065 append_json_item_on_log_server(data, get_conv_log_filename())1066 data = {1067 "tstamp": round(finish_tstamp, 4),1068 "type": "chat",1069 "model": model_name1,1070 "gen_params": {},1071 "start": round(start_tstamp, 4),1072 "finish": round(finish_tstamp, 4),1073 "state": state1.dict(),1074 "ip": get_ip(request),1075 }1076 fout.write(json.dumps(data) + "\n")1077 append_json_item_on_log_server(data, get_conv_log_filename())1078 1079 for i, state in enumerate([state0, state1]):1080 src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'1081 os.makedirs(os.path.dirname(src_img_file), exist_ok=True)1082 with open(src_img_file, 'w') as f:1083 save_any_image(state.source_image, f)1084 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'1085 with open(output_file, 'w') as f:1086 save_any_image(state.output, f)1087 save_image_file_on_log_server(src_img_file)1088 save_image_file_on_log_server(output_file)1089 1090def generate_iem_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):1091 if not model_name0:1092 raise gr.Warning("Model name A cannot be empty.")1093 if not model_name1:1094 raise gr.Warning("Model name B cannot be empty.")1095 state0 = ImageStateIE(model_name0)1096 state1 = ImageStateIE(model_name1)1097 ip = get_ip(request)1098 igm_logger.info(f"generate. ip: {ip}")1099 start_tstamp = time.time()1100 model_name0 = re.sub(r"### Model A: ", "", model_name0)1101 model_name1 = re.sub(r"### Model B: ", "", model_name1)1102 source_image, generated_image0, generated_image1, source_text, target_text, instruct_text = gen_func(model_name0, model_name1)1103 state0.source_prompt = source_text1104 state0.target_prompt = target_text1105 state0.instruct_prompt = instruct_text1106 state0.source_image = source_image1107 state0.output = generated_image01108 state0.model_name = model_name01109 state1.source_prompt = source_text1110 state1.target_prompt = target_text1111 state1.instruct_prompt = instruct_text1112 state1.source_image = source_image1113 state1.output = generated_image11114 state1.model_name = model_name11115 1116 yield state0, state1, generated_image0, generated_image1, source_image, source_text, target_text, instruct_text1117 1118 finish_tstamp = time.time()1119 # logger.info(f"===output===: {output}")1120 1121 with open(get_conv_log_filename(), "a") as fout:1122 data = {1123 "tstamp": round(finish_tstamp, 4),1124 "type": "chat",1125 "model": model_name0,1126 "gen_params": {},1127 "start": round(start_tstamp, 4),1128 "finish": round(finish_tstamp, 4),1129 "state": state0.dict(),1130 "ip": get_ip(request),1131 }1132 fout.write(json.dumps(data) + "\n")1133 append_json_item_on_log_server(data, get_conv_log_filename())1134 data = {1135 "tstamp": round(finish_tstamp, 4),1136 "type": "chat",1137 "model": model_name1,1138 "gen_params": {},1139 "start": round(start_tstamp, 4),1140 "finish": round(finish_tstamp, 4),1141 "state": state1.dict(),1142 "ip": get_ip(request),1143 }1144 fout.write(json.dumps(data) + "\n")1145 append_json_item_on_log_server(data, get_conv_log_filename())1146 1147 for i, state in enumerate([state0, state1]):1148 src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'1149 os.makedirs(os.path.dirname(src_img_file), exist_ok=True)1150 with open(src_img_file, 'w') as f:1151 save_any_image(state.source_image, f)1152 output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'1153 with open(output_file, 'w') as f:1154 save_any_image(state.output, f)1155 save_image_file_on_log_server(src_img_file)1156 save_image_file_on_log_server(output_file)1157 1158 1159def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):1160 if not source_text:1161 raise gr.Warning("Source prompt cannot be empty.")1162 if not target_text:1163 raise gr.Warning("Target prompt cannot be empty.")1164 if not instruct_text:1165 raise gr.Warning("Instruction prompt cannot be empty.")1166 if not source_image:1167 raise gr.Warning("Source image cannot be empty.")1168 state0 = ImageStateIE(model_name0)1169 state1 = ImageStateIE(model_name1)1170 ip = get_ip(request)1171 igm_logger.info(f"generate. ip: {ip}")1172 start_tstamp = time.time()1173 model_name0 = ""1174 model_name1 = ""1175 generated_image0, generated_image1, model_name0, model_name1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)1176 state0.source_prompt = source_text1177 state0.target_prompt = target_text1178 state0.instruct_prompt = instruct_text1179 state0.source_image = source_image1180 state0.output = generated_image01181 state0.model_name = model_name01182 state1.source_prompt = source_text1183 state1.target_prompt = target_text1184 state1.instruct_prompt = instruct_text1185 state1.source_image = source_image1186 state1.output = generated_image11187 state1.model_name = model_name11188 if generated_image0 == '' and generated_image1 == '':1189 with open(get_nsfw_conv_log_filename(), "a") as fout:1190 data = {1191 "type": "chat",1192 "model": model_name0,1193 "gen_params": {},1194 "start": round(start_tstamp, 4),1195 "state": state0.dict(),1196 "ip": get_ip(request),1197 }1198 fout.write(json.dumps(data) + "\n")1199 append_json_item_on_log_server(data, get_nsfw_conv_log_filename())1200 data = {