add image_cache

.gitignore

@@ -172,7 +172,8 @@ ksort-logs/
 
 *.mp4
 cache_video/
+cache_image/
 
-*generate_image_cache.py
-*generate_video_cache.py
-*get_webvid_prompt.py
+/model/models/generate_image_cache.py
+/model/models/generate_video_cache.py
+/get_webvid_prompt.py

README.md

@@ -31,14 +31,4 @@ pip install -r requirements.txt
 ## Start Hugging Face UI
 ```bash
 python app.py
-```
-
-## Start Log server
-```bash
-uvicorn serve.log_server:app --reload --port 22005 --host 0.0.0.0
-```
-
-## Update leaderboard
-```bash
-cd arena_elo && bash update_leaderboard.sh
 ```

model/fetch_museum_results/__init__.py

@@ -1,62 +0,0 @@
-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, model_name3, model_name4):
-    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
-    model_name3 = ARENA_TO_IG_MUSEUM[model_name3] if model_name3 in ARENA_TO_IG_MUSEUM else model_name3
-    model_name4 = ARENA_TO_IG_MUSEUM[model_name4] if model_name4 in ARENA_TO_IG_MUSEUM else model_name4
-
-    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
-    image_link_3 = baselink + "/" + model_name3 + "/" + uid
-    image_link_4 = baselink + "/" + model_name4 + "/" + uid
-
-    if task == "t2i": # Image Gen
-        prompt = value['prompt']
-        return [[image_link_1, image_link_2, image_link_3, image_link_4], [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, image_link_3, image_link_4], [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

@@ -1,129 +0,0 @@
-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/matchmaker.py

@@ -81,9 +81,8 @@ def load_json_via_sftp():
 class RunningPivot(object):
     running_pivot = []
 
-not_run = [12,13,14,15,16,17,18,19,20,21,22, 25,26] #23,24,
 
-def matchmaker(num_players, k_group=4):
+def matchmaker(num_players, k_group=4, not_run=[]):
     trueskill_env = TrueSkill()
     
     ratings, comparison_counts, total_comparisons = load_json_via_sftp()

model/model_manager.py

@@ -9,7 +9,7 @@ from PIL import Image
 from openai import OpenAI
 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
-from serve.upload import get_random_mscoco_prompt, get_random_video_prompt, get_ssh_random_video_prompt
+from serve.upload import get_random_mscoco_prompt, get_random_video_prompt, get_ssh_random_video_prompt, get_ssh_random_image_prompt
 from serve.constants import SSH_CACHE_OPENSOURCE, SSH_CACHE_ADVANCE, SSH_CACHE_PIKA, SSH_CACHE_SORA, SSH_CACHE_IMAGE
 
 
@@ -73,14 +73,9 @@ class ModelManager:
 
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B, model_C, model_D):
         if model_A == "" and model_B == "" and model_C == "" and model_D == "":
-            # not_run = [11, 12, 13, 14, 15, 16, 17, 18, 19]
-            # filtered_models = [model for i, model in enumerate(self.model_ig_list) if i not in not_run]
-            # model_names = random.sample([model for model in filtered_models], 4)
-
-            # model_names = random.sample([model for model in self.model_ig_list], 4)
-            
             from .matchmaker import matchmaker
-            model_ids = matchmaker(num_players=len(self.model_ig_list))
+            not_run = [12,13,14,15,16,17,18,19,20,21,22, 25,26] #23,24,
+            model_ids = matchmaker(num_players=len(self.model_ig_list), not_run=not_run)
             print(model_ids)
             model_names = [self.model_ig_list[i] for i in model_ids]
             print(model_names)
@@ -97,6 +92,26 @@ class ModelManager:
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3]
 
+    def generate_image_ig_cache_anony(self, model_A, model_B, model_C, model_D):
+        if model_A == "" and model_B == "" and model_C == "" and model_D == "":
+            from .matchmaker import matchmaker
+            not_run = [20,21,22]
+            model_ids = matchmaker(num_players=len(self.model_ig_list), not_run=not_run)
+            print(model_ids)
+            model_names = [self.model_ig_list[i] for i in model_ids]
+            print(model_names)
+        else:
+            model_names = [model_A, model_B, model_C, model_D]
+
+        root_dir = SSH_CACHE_IMAGE
+        local_dir = "./cache_image"
+        if not os.path.exists(local_dir):
+            os.makedirs(local_dir)
+        prompt, results = get_ssh_random_image_prompt(root_dir, local_dir, model_names)
+
+        return results[0], results[1], results[2], results[3], \
+            model_names[0], model_names[1], model_names[2], model_names[3], prompt
+
     def generate_video_vg_parallel_anony(self, model_A, model_B, model_C, model_D):
         if model_A == "" and model_B == "" and model_C == "" and model_D == "":
             # model_names = random.sample([model for model in self.model_vg_list], 4)
@@ -123,23 +138,6 @@ class ModelManager:
             os.makedirs(local_dir)
         prompt, results = get_ssh_random_video_prompt(root_dir, local_dir, model_names)
         cache_dir = local_dir
-        # cache_dir, prompt = get_random_video_prompt(root_dir)
-        # results = []
-        # for name in model_names:
-        #     model_source, model_name, model_type = name.split("_")
-        #     # if model_name in ["Runway-Gen3", "Pika-beta", "Pika-v1.0"]:
-        #     #     file_name = cache_dir.split("/")[-1]
-        #     #     video_path = os.path.join(cache_dir, f'{file_name}.mp4')
-        #     # else:
-        #     #     video_path = os.path.join(cache_dir, f'{model_name}.mp4')
-        #     video_path = os.path.join(cache_dir, f'{model_name}.mp4')
-        #     print(video_path)
-        #     results.append(video_path)
-
-        # with concurrent.futures.ThreadPoolExecutor() as executor:
-        #     futures = [executor.submit(self.generate_image_ig, prompt, model) if model.startswith("huggingface")
-        #                else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names]
-        #     results = [future.result() for future in futures]
 
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3], prompt, cache_dir

model/models/__init__.py

@@ -1,7 +1,3 @@
-# from .imagenhub_models import load_imagenhub_model
-# from .playground_api import load_playground_model
-# from .fal_api_models import load_fal_model
-# from .videogenhub_models import load_videogenhub_model
 from .huggingface_models import load_huggingface_model
 from .replicate_api_models import load_replicate_model
 from .openai_api_models import load_openai_model

model/models/fal_api_models.py

@@ -1,103 +0,0 @@
-import fal_client
-from PIL import Image
-import requests
-import io
-import os
-import base64
-
-FAL_MODEl_NAME_MAP = {"SDXL": "fast-sdxl", "SDXLTurbo": "fast-turbo-diffusion", "SDXLLightning": "fast-lightning-sdxl",
-                      "LCM(v1.5/XL)": "fast-lcm-diffusion", "PixArtSigma": "pixart-sigma", "StableCascade": "stable-cascade"}
-
-class FalModel():
-    def __init__(self, model_name, model_type):
-        self.model_name = model_name
-        self.model_type = model_type
-        os.environ['FAL_KEY'] = os.environ['FalAPI']
-    
-    def __call__(self, *args, **kwargs):
-        def decode_data_url(data_url):
-            # Find the start of the Base64 encoded data
-            base64_start = data_url.find(",") + 1
-            if base64_start == 0:
-                raise ValueError("Invalid data URL provided")
-
-            # Extract the Base64 encoded data
-            base64_string = data_url[base64_start:]
-
-            # Decode the Base64 string
-            decoded_bytes = base64.b64decode(base64_string)
-
-            return decoded_bytes
-
-        if self.model_type == "text2image":
-            assert "prompt" in kwargs, "prompt is required for text2image model"
-            handler = fal_client.submit(
-                f"fal-ai/{FAL_MODEl_NAME_MAP[self.model_name]}",
-                arguments={
-                    "prompt": kwargs["prompt"]
-                },
-            )
-            for event in handler.iter_events(with_logs=True):
-                if isinstance(event, fal_client.InProgress):
-                    print('Request in progress')
-                    print(event.logs)
-            result = handler.get()
-            print(result)
-            result_url = result['images'][0]['url']
-            if self.model_name in ["SDXLTurbo", "LCM(v1.5/XL)"]:
-                result_url = io.BytesIO(decode_data_url(result_url))
-                result = Image.open(result_url)
-            else:
-                response = requests.get(result_url)
-                result = Image.open(io.BytesIO(response.content))
-            return result
-        elif self.model_type == "image2image":
-            raise NotImplementedError("image2image model is not implemented yet")
-            # assert "image" in kwargs or "image_url" in kwargs, "image or image_url is required for image2image model"
-            # if "image" in kwargs:
-            #     image_url = None
-            #     pass
-            # handler = fal_client.submit(
-            #     f"fal-ai/{self.model_name}",
-            #     arguments={
-            #         "image_url": image_url
-            #     },
-            # )
-            #
-            # for event in handler.iter_events():
-            #     if isinstance(event, fal_client.InProgress):
-            #         print('Request in progress')
-            #         print(event.logs)
-            #
-            # result = handler.get()
-            # return result
-        elif self.model_type == "text2video":
-            assert "prompt" in kwargs, "prompt is required for text2video model"
-            if self.model_name == 'AnimateDiff':
-                fal_model_name = 'fast-animatediff/text-to-video'
-            elif self.model_name == 'AnimateDiffTurbo':
-                fal_model_name = 'fast-animatediff/turbo/text-to-video'
-            else:
-                raise NotImplementedError(f"text2video model of {self.model_name} in fal is not implemented yet")
-            handler = fal_client.submit(
-                f"fal-ai/{fal_model_name}",
-                arguments={
-                    "prompt": kwargs["prompt"]
-                },
-            )
-
-            for event in handler.iter_events(with_logs=True):
-                if isinstance(event, fal_client.InProgress):
-                    print('Request in progress')
-                    print(event.logs)
-
-            result = handler.get()
-            print("result video: ====")
-            print(result)
-            result_url = result['video']['url']
-            return result_url
-        else:
-            raise ValueError("model_type must be text2image or image2image")
-
-def load_fal_model(model_name, model_type):
-    return FalModel(model_name, model_type)

model/models/imagenhub_models.py

@@ -1,50 +0,0 @@
-import imagen_hub
-
-class ImagenHubModel():
-    def __init__(self, model_name):
-        self.model = imagen_hub.load(model_name)
-    
-    def __call__(self, *args, **kwargs):
-        return self.model.infer_one_image(*args, **kwargs)
-    
-class PNP(ImagenHubModel):
-    def __init__(self):
-        super().__init__('PNP')
-    
-    def __call__(self, *args, **kwargs):
-        if "num_inversion_steps" not in kwargs:
-            kwargs["num_inversion_steps"] = 200
-        return super().__call__(*args, **kwargs)
-
-class Prompt2prompt(ImagenHubModel):
-    def __init__(self):
-        super().__init__('Prompt2prompt')
-    
-    def __call__(self, *args, **kwargs):
-        if "num_inner_steps" not in kwargs:
-            kwargs["num_inner_steps"] = 3
-        return super().__call__(*args, **kwargs)
-    
-def load_imagenhub_model(model_name, model_type=None):
-    if model_name == 'PNP':
-        return PNP()
-    if model_name == 'Prompt2prompt':
-        return Prompt2prompt()
-    return ImagenHubModel(model_name)
-
-
-# for name in ['DeepFloydIF', 'PixArtAlpha', 'Kandinsky']: #, 'OpenJourney', 'LCM', 'SD' 'SDXL'
-#     # 
-#     pipe = ImagenHubModel(name)
-#     result = pipe(prompt='a cute dog is playing a ball')
-#     print(result)
-
-# for name in ['SD']:
-#     from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
-#     import torch
-#     pipe = DiffusionPipeline.from_pretrained(
-#             "stabilityai/stable-diffusion-2-base",
-#             torch_dtype=torch.float16,
-#             safety_checker=None,
-#         ).to("cuda")
-#     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

model/models/playground_api.py

@@ -1,35 +0,0 @@
-import os
-import json
-import requests
-from PIL import Image
-import io
-import base64
-class PlayGround():
-    def __init__(self, model_name, model_type=None):
-        self.model_name = model_name
-        self.model_type = model_type
-        self.api_key = os.environ['PlaygroundAPI']
-        if model_name == "PlayGroundV2":
-            self._model_name = "Playground_v2"
-        elif model_name == "PlayGroundV2.5":
-            self._model_name = "Playground_v2.5"
-
-        
-    def __call__(self, prompt):
-        headers = {
-            'Content-Type': 'application/json',
-            'Authorization': "Bearer " + self.api_key,
-        }
-
-        data = json.dumps({"prompt": prompt, "filter_model": self._model_name, "scheduler": "DPMPP_2M_K", "guidance_scale": 3})
-
-        response = requests.post('https://playground.com/api/models/external/v1', headers=headers, data=data)
-        response.raise_for_status() 
-        json_obj = response.json()
-        image_base64 = json_obj['images'][0]
-        img = Image.open(io.BytesIO(base64.decodebytes(bytes(image_base64, "utf-8"))))
-
-        return img
-    
-def load_playground_model(model_name, model_type="generation"):
-    return PlayGround(model_name, model_type)

serve/gradio_web.py

@@ -16,6 +16,7 @@ from .vote_utils import (
     generate_igm_museum,
     generate_igm_annoy,
     generate_igm_annoy_museum,
+    generate_igm_cache_annoy,
     share_js
 )
 from .Ksort import (
@@ -68,7 +69,9 @@ def build_side_by_side_ui_anony(models):
     state3 = 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)
+    gen_cache_func = partial(generate_igm_cache_annoy, models.generate_image_ig_cache_anony)
+    
+    # gen_func_random = partial(generate_igm_annoy_museum, models.generate_image_ig_museum_parallel_anony)
     
     gr.Markdown(notice_markdown, elem_id="notice_markdown")
 
@@ -216,6 +219,8 @@ def build_side_by_side_ui_anony(models):
         # draw_btn = gr.Button(value="🎲 Random sample", variant="primary", scale=0, elem_id="btnblue", elem_classes="send-button")
         send_btn = gr.Button(value="Send", variant="primary", scale=0, elem_id="btnblue")
         draw_btn = gr.Button(value="🎲 Random Prompt", variant="primary", scale=0, elem_id="btnblue")
+    with gr.Row():
+        cache_btn = gr.Button(value="🎲 Random Sample", interactive=True)
     with gr.Row():
         clear_btn = gr.Button(value="🎲 New Round", interactive=False)
         # regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
@@ -238,7 +243,7 @@ def build_side_by_side_ui_anony(models):
                   ["The bathroom with green tile and a red shower curtain.", os.path.join("./examples", "example4.jpg")]],
         inputs = [textbox, dummy_img_output])
 
-    order_btn_list = [textbox, send_btn, draw_btn, clear_btn]
+    order_btn_list = [textbox, send_btn, draw_btn, cache_btn, clear_btn]
     vote_order_list = [leftvote_btn, left1vote_btn, rightvote_btn, right1vote_btn, tie_btn, \
                        A1_btn, A2_btn, A3_btn, A4_btn, B1_btn, B2_btn, B3_btn, B4_btn, C1_btn, C2_btn, C3_btn, C4_btn, D1_btn, D2_btn, D3_btn, D4_btn, \
                        vote_textbox, vote_submit_btn, vote_mode_btn]
@@ -281,6 +286,7 @@ def build_side_by_side_ui_anony(models):
     # Top4_btn.click(reset_level, inputs=[Top4_text], outputs=[vote_level])
     vote_mode = gr.Textbox(value="Rank", visible=False, interactive=False)
     right_vote_text = gr.Textbox(value="wrong", visible=False, interactive=False)
+    cache_mode = gr.Textbox(value="True", visible=False, interactive=False)
 
     textbox.submit(
         disable_order_buttons,
@@ -314,21 +320,22 @@ def build_side_by_side_ui_anony(models):
         outputs=vote_order_list
     )
 
-    # draw_btn.click(
-    #     gen_func_random,
-    #     inputs=[state0, state1, state2, state3, model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
-    #     outputs=[state0, state1, state2, state3, generate_ig0, generate_ig1, generate_ig2, generate_ig3, chatbot_left, chatbot_left1, chatbot_right, chatbot_right1, \
-    #              textbox, model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
-    #     api_name="draw_btn_annony"
-    # ).then(
-    #     disable_order_buttons,
-    #     inputs=None,
-    #     outputs=order_btn_list
-    # ).then(
-    #     enable_vote_mode_buttons,
-    #     inputs=[vote_mode],
-    #     outputs=vote_order_list
-    # )
+    cache_btn.click(
+        disable_order_buttons,
+        inputs=[textbox, cache_mode],
+        outputs=order_btn_list
+    ).then(
+        gen_cache_func,
+        inputs=[state0, state1, state2, state3, model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
+        outputs=[state0, state1, state2, state3, generate_ig0, generate_ig1, generate_ig2, generate_ig3, chatbot_left, chatbot_left1, chatbot_right, chatbot_right1, \
+                 model_selector_left, model_selector_left1, model_selector_right, model_selector_right1, textbox],
+        api_name="send_btn_annony"
+    ).then(
+        enable_vote_mode_buttons,
+        inputs=[vote_mode, textbox],
+        outputs=vote_order_list
+    )
+
     draw_btn.click(
         get_random_mscoco_prompt,
         inputs=None, 

serve/gradio_web_video.py

@@ -310,7 +310,7 @@ def build_side_by_side_video_ui_anony(models):
                  model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
         api_name="clear_btn_annony"
     ).then(
-        enable_order_buttons,
+        enable_video_order_buttons,
         inputs=None,
         outputs=order_btn_list
     ).then(

serve/upload.py

@@ -146,8 +146,64 @@ def get_ssh_random_video_prompt(root_dir, local_dir, model_names):
     except Exception as e:
         print(f"An error occurred: {e}")
         raise NotImplementedError
+    sftp.close()
+    ssh.close()
     return prompt, local_path[1:]
 
+def get_ssh_random_image_prompt(root_dir, local_dir, model_names):
+    def is_directory(sftp, path):
+        try:
+            return stat.S_ISDIR(sftp.stat(path).st_mode)
+        except IOError:
+            return False
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    try:
+        ssh.connect(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+        sftp = ssh.open_sftp()
+
+        remote_subdirs = sftp.listdir(root_dir)
+        remote_subdirs = [d for d in remote_subdirs if is_directory(sftp, os.path.join(root_dir, d))]
+        
+        if not remote_subdirs:
+            print(f"No subdirectories found in {root_dir}")
+            raise NotImplementedError
+
+        chosen_subdir = random.choice(remote_subdirs)
+        chosen_subdir_path = os.path.join(root_dir, chosen_subdir)
+        print(f"Chosen subdirectory: {chosen_subdir_path}")
+        
+        prompt_path = 'prompt.txt'
+        results = [prompt_path]
+        for name in model_names:
+            model_source, model_name, model_type = name.split("_")
+            image_path = f'{model_name}.jpg'
+            print(image_path)
+            results.append(image_path)
+
+        local_path = []
+        for tar_file in results:
+            remote_file_path = os.path.join(chosen_subdir_path, tar_file)
+            local_file_path = os.path.join(local_dir, tar_file)
+            sftp.get(remote_file_path, local_file_path)
+            local_path.append(local_file_path)
+            print(f"Downloaded {remote_file_path} to {local_file_path}")
+
+        if os.path.exists(local_path[0]):
+            str_list = []
+            with open(local_path[0], 'r', encoding='utf-8') as file:
+                for line in file:
+                    str_list.append(line.strip())
+            prompt = str_list[0]
+        else:
+            raise NotImplementedError
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        raise NotImplementedError
+    sftp.close()
+    ssh.close()
+    return prompt, [Image.open(path) for path in local_path[1:]]
+
 def create_remote_directory(remote_directory, video=False):
     global ssh_client
     if not is_connected():

serve/utils.py

@@ -203,12 +203,17 @@ def disable_vote_mode_buttons():
 
     
 def enable_order_buttons():
-    return tuple(gr.update(interactive=True) for _ in range(4))
-def disable_order_buttons(textbox, video=False):
+    return tuple(gr.update(interactive=True) for _ in range(5))
+def disable_order_buttons(textbox, cache="False"):
+    if cache=="True":
+        return (gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True))
     if not textbox.strip():
-        return (gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True))
+        return (gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True))
     else:
-        return (gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True))
+        return (gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True))
+
+def enable_video_order_buttons():
+    return tuple(gr.update(interactive=True) for _ in range(4))
 def disable_video_order_buttons():
     return (gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True))
 

serve/vote_utils.py

@@ -871,6 +871,46 @@ def generate_igm_annoy(gen_func, state0, state1, state2, state3, text, model_nam
     #         save_any_image(state.output, f)
     #     save_image_file_on_log_server(output_file)
 
+def generate_igm_cache_annoy(gen_func, state0, state1, state2, state3, model_name0, model_name1, model_name2, model_name3, request: gr.Request):
+    if state0 is None:
+        state0 = ImageStateIG(model_name0)
+    if state1 is None:
+        state1 = ImageStateIG(model_name1)
+    if state2 is None:
+        state2 = ImageStateIG(model_name2)
+    if state3 is None:
+        state3 = ImageStateIG(model_name3)
+
+    ip = get_ip(request)
+    igm_logger.info(f"generate. ip: {ip}")
+    start_tstamp = time.time()
+    model_name0 = ""
+    model_name1 = ""
+    model_name2 = ""
+    model_name3 = ""
+
+    generated_image0, generated_image1, generated_image2, generated_image3, model_name0, model_name1, model_name2, model_name3, text \
+        = gen_func(model_name0, model_name1, model_name2, model_name3)
+    state0.prompt = text
+    state1.prompt = text
+    state2.prompt = text
+    state3.prompt = text
+
+    state0.output = generated_image0
+    state1.output = generated_image1
+    state2.output = generated_image2
+    state3.output = generated_image3
+
+    state0.model_name = model_name0
+    state1.model_name = model_name1
+    state2.model_name = model_name2
+    state3.model_name = model_name3
+
+    yield state0, state1, state2, state3, generated_image0, generated_image1, generated_image2, generated_image3, \
+        generated_image0, generated_image1, generated_image2, generated_image3, \
+        gr.Markdown(f"### Model A: {model_name0}", visible=False), gr.Markdown(f"### Model B: {model_name1}", visible=False), \
+        gr.Markdown(f"### Model C: {model_name2}", visible=False), gr.Markdown(f"### Model D: {model_name3}", visible=False), text
+
 def generate_vg_annoy(gen_func, state0, state1, state2, state3, model_name0, model_name1, model_name2, model_name3, request: gr.Request):
 
     if state0 is None: