GenAI-Arena / serve /vote_utils.py
DongfuJiang's picture
update
bd0bce1
raw
history blame
No virus
26.4 kB
import datetime
import time
import json
import uuid
import gradio as gr
from pathlib import Path
from .utils import *
from .log_utils import build_logger
from .constants import IMAGE_DIR
ig_logger = build_logger("gradio_web_server_image_generation", "gr_web_image_generation.log") # ig = image generation, loggers for single model direct chat
igm_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 battle
ie_logger = build_logger("gradio_web_server_image_editing", "gr_web_image_editing.log") # ie = image editing, loggers for single model direct chat
iem_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 battle
def vote_last_response_ig(state, vote_type, model_selector, request: gr.Request):
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(time.time(), 4),
"type": vote_type,
"model": model_selector,
"state": state.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
def vote_last_response_igm(states, vote_type, model_selectors, request: gr.Request):
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(time.time(), 4),
"type": vote_type,
"models": [x for x in model_selectors],
"states": [x.dict() for x in states],
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for state in states:
output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(output_file)
def vote_last_response_ie(state, vote_type, model_selector, request: gr.Request):
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(time.time(), 4),
"type": vote_type,
"model": model_selector,
"state": state.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
output_file = f'{IMAGE_DIR}/edition/{state.conv_id}.jpg'
source_file = f'{IMAGE_DIR}/edition/{state.conv_id}_source.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
with open(source_file, 'w') as sf:
state.source_image.save(sf, 'JPEG')
save_image_file_on_log_server(output_file)
save_image_file_on_log_server(source_file)
def vote_last_response_iem(states, vote_type, model_selectors, request: gr.Request):
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(time.time(), 4),
"type": vote_type,
"models": [x for x in model_selectors],
"states": [x.dict() for x in states],
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for state in states:
output_file = f'{IMAGE_DIR}/edition/{state.conv_id}.jpg'
source_file = f'{IMAGE_DIR}/edition/{state.conv_id}_source.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
with open(source_file, 'w') as sf:
state.source_image.save(sf, 'JPEG')
save_image_file_on_log_server(output_file)
save_image_file_on_log_server(source_file)
## Image Generation (IG) Single Model Direct Chat
def upvote_last_response_ig(state, model_selector, request: gr.Request):
ip = get_ip(request)
ig_logger.info(f"upvote. ip: {ip}")
vote_last_response_ig(state, "upvote", model_selector, request)
return ("",) + (disable_btn,) * 3
def downvote_last_response_ig(state, model_selector, request: gr.Request):
ip = get_ip(request)
ig_logger.info(f"downvote. ip: {ip}")
vote_last_response_ig(state, "downvote", model_selector, request)
return ("",) + (disable_btn,) * 3
def flag_last_response_ig(state, model_selector, request: gr.Request):
ip = get_ip(request)
ig_logger.info(f"flag. ip: {ip}")
vote_last_response_ig(state, "flag", model_selector, request)
return ("",) + (disable_btn,) * 3
## Image Generation Multi (IGM) Side-by-Side and Battle
def leftvote_last_response_igm(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
igm_logger.info(f"leftvote (named). ip: {get_ip(request)}")
vote_last_response_igm(
[state0, state1], "leftvote", [model_selector0, model_selector1], request
)
return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
def rightvote_last_response_igm(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
igm_logger.info(f"rightvote (named). ip: {get_ip(request)}")
vote_last_response_igm(
[state0, state1], "rightvote", [model_selector0, model_selector1], request
)
return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
def tievote_last_response_igm(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
igm_logger.info(f"tievote (named). ip: {get_ip(request)}")
vote_last_response_igm(
[state0, state1], "tievote", [model_selector0, model_selector1], request
)
return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
def bothbad_vote_last_response_igm(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
igm_logger.info(f"bothbad_vote (named). ip: {get_ip(request)}")
vote_last_response_igm(
[state0, state1], "bothbad_vote", [model_selector0, model_selector1], request
)
return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
## Image Editing (IE) Single Model Direct Chat
def upvote_last_response_ie(state, model_selector, request: gr.Request):
ip = get_ip(request)
ie_logger.info(f"upvote. ip: {ip}")
vote_last_response_ie(state, "upvote", model_selector, request)
return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3
def downvote_last_response_ie(state, model_selector, request: gr.Request):
ip = get_ip(request)
ie_logger.info(f"downvote. ip: {ip}")
vote_last_response_ie(state, "downvote", model_selector, request)
return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3
def flag_last_response_ie(state, model_selector, request: gr.Request):
ip = get_ip(request)
ie_logger.info(f"flag. ip: {ip}")
vote_last_response_ie(state, "flag", model_selector, request)
return ("", "", gr.Image(height=512, width=512, type="pil"), "",) + (disable_btn,) * 3
## Image Editing Multi (IEM) Side-by-Side and Battle
def leftvote_last_response_iem(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
iem_logger.info(f"leftvote (anony). ip: {get_ip(request)}")
vote_last_response_iem(
[state0, state1], "leftvote", [model_selector0, model_selector1], request
)
# names = (
# "### Model A: " + state0.model_name,
# "### Model B: " + state1.model_name,
# )
# names = (state0.model_name, state1.model_name)
names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
def rightvote_last_response_iem(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
iem_logger.info(f"rightvote (anony). ip: {get_ip(request)}")
vote_last_response_iem(
[state0, state1], "rightvote", [model_selector0, model_selector1], request
)
# names = (
# "### Model A: " + state0.model_name,
# "### Model B: " + state1.model_name,
# )
names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
def tievote_last_response_iem(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
iem_logger.info(f"tievote (anony). ip: {get_ip(request)}")
vote_last_response_iem(
[state0, state1], "tievote", [model_selector0, model_selector1], request
)
names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
def bothbad_vote_last_response_iem(
state0, state1, model_selector0, model_selector1, request: gr.Request
):
iem_logger.info(f"bothbad_vote (anony). ip: {get_ip(request)}")
vote_last_response_iem(
[state0, state1], "bothbad_vote", [model_selector0, model_selector1], request
)
names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
share_js = """
function (a, b, c, d) {
const captureElement = document.querySelector('#share-region-named');
html2canvas(captureElement)
.then(canvas => {
canvas.style.display = 'none'
document.body.appendChild(canvas)
return canvas
})
.then(canvas => {
const image = canvas.toDataURL('image/png')
const a = document.createElement('a')
a.setAttribute('download', 'chatbot-arena.png')
a.setAttribute('href', image)
a.click()
canvas.remove()
});
return [a, b, c, d];
}
"""
def share_click_igm(state0, state1, model_selector0, model_selector1, request: gr.Request):
igm_logger.info(f"share (anony). ip: {get_ip(request)}")
if state0 is not None and state1 is not None:
vote_last_response_igm(
[state0, state1], "share", [model_selector0, model_selector1], request
)
def share_click_iem(state0, state1, model_selector0, model_selector1, request: gr.Request):
iem_logger.info(f"share (anony). ip: {get_ip(request)}")
if state0 is not None and state1 is not None:
vote_last_response_iem(
[state0, state1], "share", [model_selector0, model_selector1], request
)
## All Generation Gradio Interface
class ImageStateIG:
def __init__(self, model_name):
self.conv_id = uuid.uuid4().hex
self.model_name = model_name
self.prompt = None
self.output = None
def dict(self):
base = {
"conv_id": self.conv_id,
"model_name": self.model_name,
"prompt": self.prompt
}
return base
class ImageStateIE:
def __init__(self, model_name):
self.conv_id = uuid.uuid4().hex
self.model_name = model_name
self.source_prompt = None
self.target_prompt = None
self.instruct_prompt = None
self.source_image = None
self.output = None
def dict(self):
base = {
"conv_id": self.conv_id,
"model_name": self.model_name,
"source_prompt": self.source_prompt,
"target_prompt": self.target_prompt,
"instruct_prompt": self.instruct_prompt
}
return base
def generate_ig(gen_func, state, text, model_name, request: gr.Request):
if not text:
raise gr.Warning("Prompt cannot be empty.")
if not model_name:
raise gr.Warning("Model name cannot be empty.")
if state is None:
state = ImageStateIG(model_name)
ip = get_ip(request)
ig_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image = gen_func(text, model_name)
state.prompt = text
state.output = generated_image
state.model_name = model_name
yield state, generated_image
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'
os.makedirs(os.path.dirname(output_file), exist_ok=True)
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(output_file)
def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):
if not text:
raise gr.Warning("Prompt cannot be empty.")
if not model_name0:
raise gr.Warning("Model name A cannot be empty.")
if not model_name1:
raise gr.Warning("Model name B cannot be empty.")
if state0 is None:
state0 = ImageStateIG(model_name0)
if state1 is None:
state1 = ImageStateIG(model_name1)
ip = get_ip(request)
igm_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image0, generated_image1 = gen_func(text, model_name0, model_name1)
state0.prompt = text
state1.prompt = text
state0.output = generated_image0
state1.output = generated_image1
state0.model_name = model_name0
state1.model_name = model_name1
yield state0, state1, generated_image0, generated_image1
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name0,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state0.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name1,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state1.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for i, state in enumerate([state0, state1]):
output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'
os.makedirs(os.path.dirname(output_file), exist_ok=True)
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(output_file)
def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):
if not text:
raise gr.Warning("Prompt cannot be empty.")
if state0 is None:
state0 = ImageStateIG(model_name0)
if state1 is None:
state1 = ImageStateIG(model_name1)
ip = get_ip(request)
igm_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image0, generated_image1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)
state0.prompt = text
state1.prompt = text
state0.output = generated_image0
state1.output = generated_image1
state0.model_name = model_name0
state1.model_name = model_name1
yield state0, state1, generated_image0, generated_image1, \
gr.Markdown(f"### Model A: {model_name0}"), gr.Markdown(f"### Model B: {model_name1}")
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name0,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state0.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name1,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state1.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for i, state in enumerate([state0, state1]):
output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'
os.makedirs(os.path.dirname(output_file), exist_ok=True)
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(output_file)
def generate_ie(gen_func, state, source_text, target_text, instruct_text, source_image, model_name, request: gr.Request):
if not source_text:
raise gr.Warning("Source prompt cannot be empty.")
if not target_text:
raise gr.Warning("Target prompt cannot be empty.")
if not instruct_text:
raise gr.Warning("Instruction prompt cannot be empty.")
if not source_image:
raise gr.Warning("Source image cannot be empty.")
if not model_name:
raise gr.Warning("Model name cannot be empty.")
if state is None:
state = ImageStateIE(model_name)
ip = get_ip(request)
ig_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image = gen_func(source_text, target_text, instruct_text, source_image, model_name)
state.source_prompt = source_text
state.target_prompt = target_text
state.instruct_prompt = instruct_text
state.source_image = source_image
state.output = generated_image
state.model_name = model_name
yield state, generated_image
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'
os.makedirs(os.path.dirname(src_img_file), exist_ok=True)
with open(src_img_file, 'w') as f:
state.source_image.save(f, 'JPEG')
output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(src_img_file)
save_image_file_on_log_server(output_file)
def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):
if not source_text:
raise gr.Warning("Source prompt cannot be empty.")
if not target_text:
raise gr.Warning("Target prompt cannot be empty.")
if not instruct_text:
raise gr.Warning("Instruction prompt cannot be empty.")
if not source_image:
raise gr.Warning("Source image cannot be empty.")
if not model_name0:
raise gr.Warning("Model name A cannot be empty.")
if not model_name1:
raise gr.Warning("Model name B cannot be empty.")
if state0 is None:
state0 = ImageStateIE(model_name0)
if state1 is None:
state1 = ImageStateIE(model_name1)
ip = get_ip(request)
igm_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image0, generated_image1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)
state0.source_prompt = source_text
state0.target_prompt = target_text
state0.instruct_prompt = instruct_text
state0.source_image = source_image
state0.output = generated_image0
state0.model_name = model_name0
state1.source_prompt = source_text
state1.target_prompt = target_text
state1.instruct_prompt = instruct_text
state1.source_image = source_image
state1.output = generated_image1
state1.model_name = model_name1
yield state0, state1, generated_image0, generated_image1
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name0,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state0.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name1,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state1.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for i, state in enumerate([state0, state1]):
src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'
os.makedirs(os.path.dirname(src_img_file), exist_ok=True)
with open(src_img_file, 'w') as f:
state.source_image.save(f, 'JPEG')
output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(src_img_file)
save_image_file_on_log_server(output_file)
def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):
if not source_text:
raise gr.Warning("Source prompt cannot be empty.")
if not target_text:
raise gr.Warning("Target prompt cannot be empty.")
if not instruct_text:
raise gr.Warning("Instruction prompt cannot be empty.")
if not source_image:
raise gr.Warning("Source image cannot be empty.")
if state0 is None:
state0 = ImageStateIE(model_name0)
if state1 is None:
state1 = ImageStateIE(model_name1)
ip = get_ip(request)
igm_logger.info(f"generate. ip: {ip}")
start_tstamp = time.time()
generated_image0, generated_image1, model_name0, model_name1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)
state0.source_prompt = source_text
state0.target_prompt = target_text
state0.instruct_prompt = instruct_text
state0.source_image = source_image
state0.output = generated_image0
state0.model_name = model_name0
state1.source_prompt = source_text
state1.target_prompt = target_text
state1.instruct_prompt = instruct_text
state1.source_image = source_image
state1.output = generated_image1
state1.model_name = model_name1
yield state0, state1, generated_image0, generated_image1, \
gr.Markdown(f"### Model A: {model_name0}"), gr.Markdown(f"### Model B: {model_name1}")
finish_tstamp = time.time()
# logger.info(f"===output===: {output}")
with open(get_conv_log_filename(), "a") as fout:
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name0,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state0.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
data = {
"tstamp": round(finish_tstamp, 4),
"type": "chat",
"model": model_name1,
"gen_params": {},
"start": round(start_tstamp, 4),
"finish": round(finish_tstamp, 4),
"state": state1.dict(),
"ip": get_ip(request),
}
fout.write(json.dumps(data) + "\n")
append_json_item_on_log_server(data, get_conv_log_filename())
for i, state in enumerate([state0, state1]):
src_img_file = f'{IMAGE_DIR}/edition/{state.conv_id}_src.jpg'
os.makedirs(os.path.dirname(src_img_file), exist_ok=True)
with open(src_img_file, 'w') as f:
state.source_image.save(f, 'JPEG')
output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
with open(output_file, 'w') as f:
state.output.save(f, 'JPEG')
save_image_file_on_log_server(src_img_file)
save_image_file_on_log_server(output_file)