fix regenerate bug

serve/constants.py

@@ -17,4 +17,16 @@ SAVE_IMAGE = "save_image"
 SAVE_VIDEO = "save_video"
 SAVE_LOG = "save_log"
 
+IMAGE_GENERATION_MODELS = ['fal_LCM(v1.5/XL)_text2image','fal_SDXLTurbo_text2image','fal_SDXL_text2image', 'imagenhub_PixArtAlpha_generation', 'fal_PixArtSigma_text2image',
+                            'imagenhub_OpenJourney_generation','fal_SDXLLightning_text2image', 'fal_StableCascade_text2image',
+                            'playground_PlayGroundV2_generation', 'playground_PlayGroundV2.5_generation']
+IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
+                        'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition',
+                        'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition',
+                        'imagenhub_InfEdit_edition', 'imagenhub_CosXLEdit_edition']
+VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
+                           'fal_AnimateDiffTurbo_text2video',
+                           'videogenhub_LaVie_generation', 'videogenhub_VideoCrafter2_generation',
+                           'videogenhub_ModelScope_generation', 'videogenhub_OpenSora_generation']
+
 

serve/vote_utils.py

@@ -7,7 +7,7 @@ import regex as re
 from pathlib import Path
 from .utils import *
 from .log_utils import build_logger
-from .constants import IMAGE_DIR, VIDEO_DIR
+from .constants import IMAGE_DIR, VIDEO_DIR, IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS
 import imageio
 
 ig_logger = build_logger("gradio_web_server_image_generation", "gr_web_image_generation.log") # ig = image generation, loggers for single model direct chat
@@ -177,7 +177,8 @@ def leftvote_last_response_igm(
         gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
         gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        return ("",) + (disable_btn,) * 4 + (gr.Markdown('', visible=True), gr.Markdown('', visible=True))
+        return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),
+                                             gr.Markdown(state1.model_name, visible=True))
 
 def rightvote_last_response_igm(
     state0, state1, model_selector0, model_selector1, request: gr.Request
@@ -190,8 +191,8 @@ def rightvote_last_response_igm(
     if model_selector0 == "":
         return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        print("no")
-        return ("",) + (disable_btn,) * 4 + (gr.Markdown('', visible=True), gr.Markdown('', visible=True))
+        return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),
+                                             gr.Markdown(state1.model_name, visible=True))
 
 
 def tievote_last_response_igm(
@@ -206,7 +207,8 @@ def tievote_last_response_igm(
         gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
         gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        return ("",) + (disable_btn,) * 4 + (gr.Markdown('', visible=True), gr.Markdown('', visible=True))
+        return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),
+                                             gr.Markdown(state1.model_name, visible=True))
 
 
 def bothbad_vote_last_response_igm(
@@ -221,7 +223,8 @@ def bothbad_vote_last_response_igm(
             gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
             gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        return ("",) + (disable_btn,) * 4 + (gr.Markdown('', visible=True), gr.Markdown('', visible=True))
+        return ("",) + (disable_btn,) * 4 + (gr.Markdown(state0.model_name, visible=True),
+                                             gr.Markdown(state1.model_name, visible=True))
 
 ## Image Editing (IE) Single Model Direct Chat
 
@@ -259,7 +262,7 @@ def leftvote_last_response_iem(
     if model_selector0 == "":
         names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        names = (gr.Markdown('', visible=False), gr.Markdown('', visible=False))
+        names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))
     return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
         
 def rightvote_last_response_iem(
@@ -277,7 +280,7 @@ def rightvote_last_response_iem(
         names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
                  gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        names = (gr.Markdown('', visible=False), gr.Markdown('', visible=False))
+        names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))
     return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
 
 def tievote_last_response_iem(
@@ -291,7 +294,7 @@ def tievote_last_response_iem(
         names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
                  gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        names = (gr.Markdown('', visible=False), gr.Markdown('', visible=False))
+        names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))
     return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
 
 def bothbad_vote_last_response_iem(
@@ -305,7 +308,7 @@ def bothbad_vote_last_response_iem(
         names = (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True),
                  gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
-        names = (gr.Markdown('', visible=False), gr.Markdown('', visible=False))
+        names = (gr.Markdown(state0.model_name, visible=False), gr.Markdown(state1.model_name, visible=False))
     return names + ("", "", gr.Image(height=512, width=512, type="pil"), "") + (disable_btn,) * 4
 
 
@@ -342,8 +345,8 @@ def leftvote_last_response_vgm(
         return ("",) + (disable_btn,) * 4 + (gr.Markdown(f"### Model A: {state0.model_name.split('_')[1]}", visible=True), gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
         return ("",) + (disable_btn,) * 4 + (
-        gr.Markdown("", visible=False),
-        gr.Markdown("", visible=False))
+        gr.Markdown(state0.model_name, visible=False),
+        gr.Markdown(state1.model_name, visible=False))
 
 
 def rightvote_last_response_vgm(
@@ -359,8 +362,8 @@ def rightvote_last_response_vgm(
         gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
         return ("",) + (disable_btn,) * 4 + (
-            gr.Markdown("", visible=False),
-            gr.Markdown("", visible=False))
+            gr.Markdown(state0.model_name, visible=False),
+            gr.Markdown(state1.model_name, visible=False))
 
 def tievote_last_response_vgm(
     state0, state1, model_selector0, model_selector1, request: gr.Request
@@ -375,8 +378,8 @@ def tievote_last_response_vgm(
         gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
         return ("",) + (disable_btn,) * 4 + (
-            gr.Markdown("", visible=False),
-            gr.Markdown("", visible=False))
+            gr.Markdown(state0.model_name, visible=False),
+            gr.Markdown(state1.model_name, visible=False))
 
 
 def bothbad_vote_last_response_vgm(
@@ -392,8 +395,8 @@ def bothbad_vote_last_response_vgm(
         gr.Markdown(f"### Model B: {state1.model_name.split('_')[1]}", visible=True))
     else:
         return ("",) + (disable_btn,) * 4 + (
-            gr.Markdown("", visible=False),
-            gr.Markdown("", visible=False))
+            gr.Markdown(state0.model_name, visible=False),
+            gr.Markdown(state1.model_name, visible=False))
 
 share_js = """
 function (a, b, c, d) {
@@ -482,8 +485,6 @@ class VideoStateVG:
         return base
 
 
-
-
 def generate_ig(gen_func, state, text, model_name, request: gr.Request):
     if not text:
         raise gr.Warning("Prompt cannot be empty.")
@@ -599,6 +600,11 @@ def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
     start_tstamp = time.time()
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    model_map = {model_name.split('_')[1]: model_name for model_name in IMAGE_GENERATION_MODELS}
+    if model_name0 in model_map:
+        model_name0 = model_map[model_name0]
+    if model_name1 in model_map:
+        model_name1 = model_map[model_name1]
     generated_image0, generated_image1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)
     state0.prompt = text
     state1.prompt = text
@@ -797,6 +803,11 @@ def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instr
     start_tstamp = time.time()
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    model_map = {model_name.split('_')[1]: model_name for model_name in IMAGE_EDITION_MODELS}
+    if model_name0 in model_map:
+        model_name0 = model_map[model_name0]
+    if model_name1 in model_map:
+        model_name1 = model_map[model_name1]
     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
@@ -991,6 +1002,11 @@ def generate_vgm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
     start_tstamp = time.time()
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
+    model_map = {model_name.split('_')[1]: model_name for model_name in VIDEO_GENERATION_MODELS}
+    if model_name0 in model_map:
+        model_name0 = model_map[model_name0]
+    if model_name1 in model_map:
+        model_name1 = model_map[model_name1]
     generated_video0, generated_video1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)
     state0.prompt = text
     state1.prompt = text