Adding port from museum

model/fetch_museum_results/__init__.py

@@ -0,0 +1,58 @@
+from .imagen_museum import TASK_DICT, DOMAIN
+from .imagen_museum import fetch_indexes
+import random
+
+ARENA_TO_IG_MUSEUM = {"LCM(v1.5/XL)":"LCM", 
+                   "PlayGroundV2.5": "PlayGroundV2_5"}
+
+def draw2_from_imagen_museum(task, model_name1, model_name2):
+    task_name = TASK_DICT[task]
+    model_name1 = ARENA_TO_IG_MUSEUM[model_name1] if model_name1 in ARENA_TO_IG_MUSEUM else model_name1
+    model_name2 = ARENA_TO_IG_MUSEUM[model_name2] if model_name2 in ARENA_TO_IG_MUSEUM else model_name2
+
+    domain = DOMAIN
+    baselink = domain + task_name
+
+    matched_results = fetch_indexes(baselink)
+    r = random.Random()
+    uid, value = r.choice(list(matched_results.items()))
+    image_link_1 = baselink + "/" + model_name1 + "/" + uid
+    image_link_2 = baselink + "/" + model_name2 + "/" + uid
+
+    if task == "t2i": # Image Gen
+        prompt = value['prompt']
+        return [[image_link_1, image_link_2], [prompt]]
+    if task == "tie": # Image Edit
+        instruction = value['instruction']
+        input_caption = value['source_global_caption']
+        output_caption = value['target_global_caption']
+        source_image_link = baselink + "/" + "input" + "/" + uid
+        return [[source_image_link, image_link_1, image_link_2], [input_caption, output_caption, instruction]]
+    else:
+        raise ValueError("Task not supported")
+
+def draw_from_imagen_museum(task, model_name):
+    task_name = TASK_DICT[task]
+    model_name = ARENA_TO_IG_MUSEUM[model_name] if model_name in ARENA_TO_IG_MUSEUM else model_name
+
+    domain = DOMAIN
+    baselink = domain + task_name
+
+    matched_results = fetch_indexes(baselink)
+    r = random.Random()
+    uid, value = r.choice(list(matched_results.items()))
+    model = model_name
+    image_link = baselink + "/" + model + "/" + uid
+    print(image_link)
+
+    if task == "t2i": # Image Gen
+        prompt = value['prompt']
+        return [image_link, prompt]
+    if task == "tie": # Image Edit
+        instruction = value['instruction']
+        input_caption = value['source_global_caption']
+        output_caption = value['target_global_caption']
+        source_image_link = baselink + "/" + "input" + "/" + uid
+        return [[source_image_link, image_link], [input_caption, output_caption, instruction]]
+    else:
+        raise ValueError("Task not supported")

model/fetch_museum_results/imagen_museum/__init__.py

@@ -0,0 +1,129 @@
+import csv
+import requests
+from io import StringIO
+from typing import Union, Optional, Tuple
+from PIL import Image
+import random
+
+__version__ = "0.0.1_GenAI_Arena"
+
+DOMAIN = "https://chromaica.github.io/Museum/"
+
+TASK_DICT = {
+    "t2i": "ImagenHub_Text-Guided_IG",
+    "tie": "ImagenHub_Text-Guided_IE",
+    "mie": "ImagenHub_Control-Guided_IG",
+    "cig": "ImagenHub_Control-Guided_IE",
+    "msdig": "ImagenHub_Multi-Concept_IC",
+    "sdig": "ImagenHub_Subject-Driven_IG",
+    "sdie": "ImagenHub_Subject-Driven_IE"
+}
+
+t2i_models= [
+      "SD",
+      "SDXL",
+      "OpenJourney",
+      "DeepFloydIF",
+      "DALLE2"
+    ]
+
+mie_models =  [
+      "Glide",
+      "SDInpaint",
+      "BlendedDiffusion",
+      "SDXLInpaint"
+    ] 
+
+tie_models = [
+      "DiffEdit",
+      "MagicBrush",
+      "InstructPix2Pix",
+      "Prompt2prompt",
+      "Text2Live",
+      "SDEdit",
+      "CycleDiffusion",
+      "Pix2PixZero"
+    ]
+
+sdig_models = [
+      "DreamBooth",
+      "DreamBoothLora",
+      "TextualInversion",
+      "BLIPDiffusion_Gen"
+    ]
+
+sdie_models = [
+      "PhotoSwap",
+      "DreamEdit",
+      "BLIPDiffusion_Edit"
+    ] 
+
+msdig_models = [
+      "DreamBooth",
+      "CustomDiffusion",
+      "TextualInversion"
+    ]
+
+cig_models = [
+      "ControlNet",
+      "UniControl"
+    ]
+
+def fetch_csv_keys(url):
+    """
+    Fetches a CSV file from a given URL and parses it into a list of keys,
+    ignoring the header line.
+    """
+    response = requests.get(url)
+    response.raise_for_status()  # Ensure we notice bad responses
+    
+    # Use StringIO to turn the fetched text data into a file-like object
+    csv_file = StringIO(response.text)
+
+    # Create a CSV reader
+    csv_reader = csv.reader(csv_file)
+
+    # Skip the header
+    next(csv_reader, None)
+
+    # Return the list of keys
+    return [row[0] for row in csv_reader if row]
+
+def fetch_json_data(url):
+    """
+    Fetches JSON data from a given URL.
+    """
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.json()
+
+def fetch_data_and_match(csv_url, json_url):
+    """
+    Fetches a list of keys from a CSV and then fetches JSON data and matches the keys to the JSON.
+    """
+    # Fetch keys from CSV
+    keys = fetch_csv_keys(csv_url)
+    
+    # Fetch JSON data
+    json_data = fetch_json_data(json_url)
+    
+    # Extract relevant data using keys
+    matched_data = {key: json_data.get(key) for key in keys if key in json_data}
+
+    return matched_data
+
+def fetch_indexes(baselink):
+    matched_results = fetch_data_and_match(baselink+"/dataset_lookup.csv", baselink+"/dataset_lookup.json")
+    return matched_results
+
+if __name__ == "__main__":
+    domain = "https://chromaica.github.io/Museum/"
+    baselink = domain + "ImagenHub_Text-Guided_IE"
+    matched_results = fetch_indexes(baselink)
+    for uid, value in matched_results.items():
+        print(uid)
+        model = "CycleDiffusion"
+        image_link = baselink + "/" + model + "/" + uid
+        print(image_link)
+        instruction = value['instruction']
+        print(instruction)

model/model_manager.py

@@ -6,6 +6,7 @@ import io, base64, json
 import spaces
 from PIL import Image
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, load_pipeline
+from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum
 
 class ModelManager:
     def __init__(self):
@@ -33,6 +34,15 @@ class ModelManager:
         result = pipe(prompt=prompt)
         return result
 
+    def generate_image_ig_museum(self, model_name):
+        model_name = model_name.split('_')[1]
+        result_list = draw_from_imagen_museum("t2i", model_name)
+        image_link = result_list[0]
+        prompt = result_list[1]
+
+        return image_link, prompt
+
+
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B):
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in self.model_ig_list], 2)
@@ -45,6 +55,21 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1], model_names[0], model_names[1]
 
+    def generate_image_ig_museum_parallel_anony(self, model_A, model_B):
+        if model_A == "" and model_B == "":
+            model_names = random.sample([model for model in self.model_ig_list], 2)
+        else:
+            model_names = [model_A, model_B]
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            model_1 = model_names[0].split('_')[1]
+            model_2 = model_names[1].split('_')[1]
+            result_list = draw2_from_imagen_museum("t2i", model_1, model_2)
+            image_links = result_list[0]
+            prompt_list = result_list[1]
+
+        return image_links[0], image_links[1], model_names[0], model_names[1], prompt_list[0]
+    
     def generate_image_ig_parallel(self, prompt, model_A, model_B):
         model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -53,12 +78,31 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
 
+    def generate_image_ig_museum_parallel(self, model_A, model_B):
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            model_1 = model_A.split('_')[1]
+            model_2 = model_B.split('_')[1]
+            result_list = draw2_from_imagen_museum("t2i", model_1, model_2)
+            image_links = result_list[0]
+            prompt_list = result_list[1]
+        return image_links[0], image_links[1], prompt_list[0]
+
+
     @spaces.GPU(duration=200)
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
         return result
 
+    def generate_image_ie_museum(self, model_name):
+        model_name = model_name.split('_')[1]
+        result_list = draw_from_imagen_museum("tie", model_name)
+        image_links = result_list[0]
+        prompt_list = result_list[1]
+        # image_links = [src, model]
+        # prompt_list = [source_caption, target_caption, instruction]
+        return image_links[0], image_links[1], prompt_list[0], prompt_list[1], prompt_list[2]
+
     def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -68,6 +112,18 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
 
+    def generate_image_ie_museum_parallel(self, model_A, model_B):
+        model_names = [model_A, model_B]
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            model_1 = model_names[0].split('_')[1]
+            model_2 = model_names[1].split('_')[1]
+            result_list = draw2_from_imagen_museum("tie", model_1, model_2)
+            image_links = result_list[0]
+            prompt_list = result_list[1]
+            # image_links = [src, model_A, model_B]
+            # prompt_list = [source_caption, target_caption, instruction]
+        return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2]
+
     def generate_image_ie_parallel_anony(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in self.model_ie_list], 2)
@@ -78,6 +134,21 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1], model_names[0], model_names[1]
 
+    def generate_image_ie_museum_parallel_anony(self, model_A, model_B):
+        if model_A == "" and model_B == "":
+            model_names = random.sample([model for model in self.model_ie_list], 2)
+        else:
+            model_names = [model_A, model_B]
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            model_1 = model_names[0].split('_')[1]
+            model_2 = model_names[1].split('_')[1]
+            result_list = draw2_from_imagen_museum("tie", model_1, model_2)
+            image_links = result_list[0]
+            prompt_list = result_list[1]
+            # image_links = [src, model_A, model_B]
+            # prompt_list = [source_caption, target_caption, instruction]
+        return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2], model_names[0], model_names[1]
+
     @spaces.GPU(duration=120)
     def generate_video_vg(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
@@ -89,6 +160,9 @@ class ModelManager:
         result = pipe(prompt=prompt)
         return result
 
+    def generate_video_vg_museum(self, model_name):
+        raise NotImplementedError
+
     def generate_video_vg_parallel_anony(self, prompt, model_A, model_B):
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in self.model_vg_list], 2)
@@ -101,10 +175,23 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1], model_names[0], model_names[1]
 
+    def generate_video_vg_museum_parallel_anony(self, model_A, model_B):
+        if model_A == "" and model_B == "":
+            model_names = random.sample([model for model in self.model_vg_list], 2)
+        else:
+            model_names = [model_A, model_B]
+
+        raise NotImplementedError
+    
     def generate_video_vg_parallel(self, prompt, model_A, model_B):
         model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
             futures = [executor.submit(self.generate_video_vg, prompt, model) if model.startswith("videogenhub")
                        else executor.submit(self.generate_video_vg_api, prompt, model) for model in model_names]
             results = [future.result() for future in futures]
-        return results[0], results[1]
+        return results[0], results[1]
+
+    def generate_video_vg_museum_parallel(self, model_A, model_B):
+        model_names = [model_A, model_B]
+
+        raise NotImplementedError

serve/gradio_web.py

@@ -9,8 +9,11 @@ from .vote_utils import (
     bothbad_vote_last_response_igm as bothbad_vote_last_response,
     share_click_igm as share_click,
     generate_ig,
+    generate_ig_museum,
     generate_igm,
+    generate_igm_museum,
     generate_igm_annoy,
+    generate_igm_annoy_museum,
     share_js
 )
 from functools import partial
@@ -37,6 +40,7 @@ Find out who is the 🥇conditional image generation models! More models are goi
     state0 = gr.State()
     state1 = gr.State()
     gen_func = partial(generate_igm_annoy, models.generate_image_ig_parallel_anony)
+    gen_func_random = partial(generate_igm_annoy_museum, models.generate_image_ig_museum_parallel_anony)
     
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
@@ -78,6 +82,7 @@ Find out who is the 🥇conditional image generation models! More models are goi
             elem_id="input_box",
         )
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row():
         clear_btn = gr.Button(value="🎲 New Round", interactive=False)
@@ -117,6 +122,16 @@ Find out who is the 🥇conditional image generation models! More models are goi
         inputs=None,
         outputs=btn_list
     )
+    draw_btn.click(
+        gen_func_random,
+        inputs=[state0, state1, model_selector_left, model_selector_right],
+        outputs=[state0, state1, chatbot_left, chatbot_right, textbox, model_selector_left, model_selector_right],
+        api_name="draw_btn_annony"
+    ).then(
+        enable_buttons_side_by_side,
+        inputs=None,
+        outputs=btn_list
+    )
     
     clear_btn.click(
         clear_history_side_by_side_anony, 
@@ -189,6 +204,7 @@ def build_side_by_side_ui_named(models):
     state1 = gr.State()
     anony = False
     gen_func = partial(generate_igm, models.generate_image_ig_parallel)
+    gen_func_random = partial(generate_igm_museum, models.generate_image_ig_museum_parallel)
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
     with gr.Group(elem_id="share-region-named"):
@@ -240,6 +256,7 @@ def build_side_by_side_ui_named(models):
             elem_id="input_box"
         )
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row():
         clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
@@ -282,6 +299,16 @@ def build_side_by_side_ui_named(models):
         inputs=None,  
         outputs=btn_list 
     )
+    draw_btn.click(
+        gen_func_random,
+        inputs=[state0, state1, model_selector_left, model_selector_right],
+        outputs=[state0, state1, chatbot_left, chatbot_right, textbox],
+        api_name="draw_side_by_side"
+    ).then(
+        enable_buttons_side_by_side,
+        inputs=None,
+        outputs=btn_list
+    )
     regenerate_btn.click(
         gen_func,
         inputs=[state0, state1, textbox, model_selector_left, model_selector_right],
@@ -350,6 +377,7 @@ def build_single_model_ui(models, add_promotion_links=False):
 
     state = gr.State()
     gen_func = partial(generate_ig, models.generate_image_ig)
+    gen_func_random = partial(generate_ig_museum, models.generate_image_ig_museum)
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
     model_list = models.model_ig_list
@@ -379,6 +407,7 @@ def build_single_model_ui(models, add_promotion_links=False):
         )
 
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row():
         chatbot = gr.Image()
@@ -431,7 +460,17 @@ def build_single_model_ui(models, add_promotion_links=False):
         inputs=None,  
         outputs=btn_list 
     )
-    
+    draw_btn.click(
+        gen_func_random, 
+        inputs=[state, model_selector], 
+        outputs=[state, chatbot, textbox],
+        api_name="draw_btn_single",
+        show_progress = "full"
+    ).success(
+        enable_buttons, 
+        inputs=None,  
+        outputs=btn_list 
+    )
     upvote_btn.click(
         upvote_last_response,
         inputs=[state, model_selector],

serve/gradio_web_image_editing.py

@@ -9,8 +9,11 @@ from .vote_utils import (
     bothbad_vote_last_response_iem as bothbad_vote_last_response,
     share_click_iem as share_click,
     generate_ie,
+    generate_ie_museum,
     generate_iem,
+    generate_iem_museum,
     generate_iem_annoy,
+    generate_iem_annoy_museum,
     share_js
 )
 from functools import partial
@@ -40,6 +43,7 @@ Find out who is the 🥇conditional image edition models!
 
     state0, state1 = gr.State(), gr.State()
     gen_func = partial(generate_iem_annoy, models.generate_image_ie_parallel_anony)
+    gen_func_random = partial(generate_iem_annoy_museum, models.generate_image_ie_museum_parallel_anony)
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
     with gr.Group(elem_id="share-region-anony"):
@@ -93,6 +97,7 @@ Find out who is the 🥇conditional image edition models!
     with gr.Row():
         source_image = gr.Image(type="pil")
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row() as button_row:
         clear_btn = gr.Button(value="🎲 New Round", interactive=False)
@@ -138,6 +143,23 @@ Find out who is the 🥇conditional image edition models!
         inputs=None,  
         outputs=btn_list 
     )
+    draw_btn.click(
+        gen_func_random,
+        inputs=[
+            state0,
+            state1,
+            model_selector_left, 
+            model_selector_right
+        ],
+        outputs=[state0, state1, chatbot_left, chatbot_right, 
+                source_image, textbox_source, textbox_target, textbox_instruct, 
+                model_selector_left, model_selector_right],
+        api_name="draw_btn_side_by_side_anony"
+    ).then(
+        enable_buttons_side_by_side, 
+        inputs=None,  
+        outputs=btn_list 
+    )
     regenerate_btn.click(
         gen_func,
         inputs=[
@@ -219,6 +241,7 @@ def build_side_by_side_ui_named_ie(models):
 
     state0, state1 = gr.State(), gr.State()
     gen_func = partial(generate_iem, models.generate_image_ie_parallel)
+    gen_func_random = partial(generate_iem_museum, models.generate_image_ie_museum_parallel)
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
     with gr.Group(elem_id="share-region-named"):
@@ -283,6 +306,7 @@ def build_side_by_side_ui_named_ie(models):
     with gr.Row():
         source_image = gr.Image(type="pil")
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row() as button_row:
         clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
@@ -334,6 +358,23 @@ def build_side_by_side_ui_named_ie(models):
         inputs=None,  
         outputs=btn_list 
     )
+    draw_btn.click(
+        gen_func_random,
+        inputs=[
+            state0,
+            state1,
+            model_selector_left, 
+            model_selector_right
+        ],
+        outputs=[state0, state1, chatbot_left, chatbot_right, 
+                 source_image, textbox_source, textbox_target, textbox_instruct 
+                ],
+        api_name="draw_btn_side_by_side"
+    ).then(
+        enable_buttons_side_by_side, 
+        inputs=None,  
+        outputs=btn_list 
+    )
     regenerate_btn.click(
         gen_func,
         inputs=[
@@ -409,6 +450,7 @@ def build_single_model_ui_ie(models, add_promotion_links=False):
 
     state = gr.State()
     gen_func = partial(generate_ie, models.generate_image_ie)
+    gen_func_random = partial(generate_ie_museum, models.generate_image_ie_museum)
     model_list = models.model_ie_list
 
     with gr.Row(elem_id="model_selector_row"):
@@ -448,6 +490,7 @@ def build_single_model_ui_ie(models, add_promotion_links=False):
     with gr.Row():
         source_image = gr.Image(type="pil")
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
+        draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0)
 
     with gr.Row():
         chatbot = gr.Image()
@@ -493,7 +536,19 @@ def build_single_model_ui_ie(models, add_promotion_links=False):
         inputs=None,  
         outputs=btn_list 
     )
+    draw_btn.click(
+        gen_func_random,
+        inputs=[state, model_selector], 
+        outputs=[state, chatbot, source_image, textbox_source, textbox_target, textbox_instruct], 
+        api_name="send_btn_single",
+        show_progress = "full"
+    ).then(
+        enable_buttons, 
+        inputs=None,  
+        outputs=btn_list 
+    )
     
+
     vote_btns = [upvote_btn, downvote_btn, flag_btn]
     model_inputs = [textbox_source, textbox_instruct, source_image, textbox_target]
     upvote_btn.click(

serve/vote_utils.py

@@ -9,6 +9,7 @@ from .utils import *
 from .log_utils import build_logger
 from .constants import IMAGE_DIR, VIDEO_DIR
 import imageio
+from diffusers.utils import load_image
 
 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
@@ -17,6 +18,13 @@ iem_logger = build_logger("gradio_web_server_image_editing_multi", "gr_web_image
 vg_logger = build_logger("gradio_web_server_video_generation", "gr_web_video_generation.log") # vg = video generation, loggers for single model direct chat
 vgm_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 battle
 
+def save_any_image(image_file, file_path):
+    if isinstance(image_file, str):
+        image = load_image(image_file)
+        image.save(file_path, 'JPEG')
+    else:
+        image_file.save(file_path, 'JPEG')
+
 def vote_last_response_ig(state, vote_type, model_selector, request: gr.Request):
     with open(get_conv_log_filename(), "a") as fout:
         data = {
@@ -30,7 +38,7 @@ def vote_last_response_ig(state, vote_type, model_selector, request: gr.Request)
         append_json_item_on_log_server(data, get_conv_log_filename())
     output_file = f'{IMAGE_DIR}/generation/{state.conv_id}.jpg'
     with open(output_file, 'w') as f:
-        state.output.save(f, 'JPEG')
+        save_any_image(state.output, f)
     save_image_file_on_log_server(output_file)
         
 def vote_last_response_igm(states, vote_type, model_selectors, request: gr.Request):
@@ -47,7 +55,7 @@ def vote_last_response_igm(states, vote_type, model_selectors, request: gr.Reque
     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_any_image(state.output, f)
         save_image_file_on_log_server(output_file)
 
 def vote_last_response_ie(state, vote_type, model_selector, request: gr.Request):
@@ -64,9 +72,9 @@ def vote_last_response_ie(state, vote_type, model_selector, request: gr.Request)
     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')
+        save_any_image(state.output, f)
     with open(source_file, 'w') as sf:
-        state.source_image.save(sf, 'JPEG')
+        save_any_image(state.source_image, sf)
     save_image_file_on_log_server(output_file)
     save_image_file_on_log_server(source_file)
         
@@ -85,9 +93,9 @@ def vote_last_response_iem(states, vote_type, model_selectors, request: gr.Reque
         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')
+            save_any_image(state.output, f)
         with open(source_file, 'w') as sf:
-            state.source_image.save(sf, 'JPEG')
+            save_any_image(state.source_image, sf)
         save_image_file_on_log_server(output_file)
         save_image_file_on_log_server(source_file)
 
@@ -522,9 +530,47 @@ def generate_ig(gen_func, state, text, model_name, request: gr.Request):
     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_any_image(state.output, f)
     save_image_file_on_log_server(output_file)
+
+def generate_ig_museum(gen_func, state, model_name, request: gr.Request):
+    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, text = gen_func(model_name)
+    state.prompt = text
+    state.output = generated_image
+    state.model_name = model_name
+    
+    yield state, generated_image, text
+    
+    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:
+        save_any_image(state.output, f)
+    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.")
@@ -585,9 +631,71 @@ def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, reque
         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_any_image(state.output, f)
         save_image_file_on_log_server(output_file)
-            
+
+def generate_igm_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):
+    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()
+    # Remove ### Model (A|B): from model name
+    model_name0 = re.sub(r"### Model A: ", "", model_name0)
+    model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    generated_image0, generated_image1, text = gen_func(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, text
+    
+    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:
+            save_any_image(state.output, f)
+        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.")
@@ -644,9 +752,65 @@ def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
         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_any_image(state.output, f)
         save_image_file_on_log_server(output_file)
             
+def generate_igm_annoy_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):
+    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()
+    model_name0 = re.sub(r"### Model A: ", "", model_name0)
+    model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    generated_image0, generated_image1, model_name0, model_name1, text = gen_func(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, text,\
+        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:
+            save_any_image(state.output, f)
+        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:
@@ -694,13 +858,59 @@ def generate_ie(gen_func, state, source_text, target_text, instruct_text, source
     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')
+        save_any_image(state.source_image, f)
     output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
     with open(output_file, 'w') as f:
-        state.output.save(f, 'JPEG')
+        save_any_image(state.output, f)
     save_image_file_on_log_server(src_img_file)
     save_image_file_on_log_server(output_file)
+
+def generate_ie_museum(gen_func, state, model_name, request: gr.Request):
+    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()
+    source_image, generated_image, source_text, target_text, instruct_text = gen_func(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, source_image, source_text, target_text, instruct_text
+    
+    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:
+        save_any_image(state.source_image, f)
+    output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
+    with open(output_file, 'w') as f:
+        save_any_image(state.output, f)
+    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.")
@@ -772,10 +982,80 @@ def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_te
         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')
+            save_any_image(state.source_image, f)
         output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
         with open(output_file, 'w') as f:
-            state.output.save(f, 'JPEG')
+            save_any_image(state.output, f)
+        save_image_file_on_log_server(src_img_file)
+        save_image_file_on_log_server(output_file)
+
+def generate_iem_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):
+    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()
+    model_name0 = re.sub(r"### Model A: ", "", model_name0)
+    model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    source_image, generated_image0, generated_image1, source_text, target_text, instruct_text = gen_func(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, source_image, source_text, target_text, instruct_text
+    
+    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:
+            save_any_image(state.source_image, f)
+        output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
+        with open(output_file, 'w') as f:
+            save_any_image(state.output, f)
         save_image_file_on_log_server(src_img_file)
         save_image_file_on_log_server(output_file)
             
@@ -848,10 +1128,77 @@ def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instr
         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')
+            save_any_image(state.source_image, f)
+        output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
+        with open(output_file, 'w') as f:
+            save_any_image(state.output, f)
+        save_image_file_on_log_server(src_img_file)
+        save_image_file_on_log_server(output_file)
+
+def generate_iem_annoy_museum(gen_func, state0, state1, model_name0, model_name1, request: gr.Request):
+    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()
+    model_name0 = ""
+    model_name1 = ""
+    source_image, generated_image0, generated_image1, source_text, target_text, instruct_text, model_name0, model_name1 = gen_func(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, source_image, source_text, target_text, instruct_text, \
+        gr.Markdown(f"### Model A: {model_name0}", visible=False), gr.Markdown(f"### Model B: {model_name1}", visible=False)
+    
+    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:
+            save_any_image(state.source_image, f)
         output_file = f'{IMAGE_DIR}/edition/{state.conv_id}_out.jpg'
         with open(output_file, 'w') as f:
-            state.output.save(f, 'JPEG')
+            save_any_image(state.output, f)
         save_image_file_on_log_server(src_img_file)
         save_image_file_on_log_server(output_file)