K-Sort-Arena

Running on Zero

App Files Files Community

YangZhoumill commited on Aug 26, 2024

Commit

a4b32da

1 Parent(s): 780769f

release code

Browse files

Files changed (14) hide show

app.py +5 -8
model/matchmaker.py +11 -5
model/matchmaker_video.py +12 -5
model/model_manager.py +0 -8
model/model_registry.py +0 -185
model/models/__init__.py +2 -18
model/models/openai_api_models.py +0 -2
model/models/other_api_models.py +3 -27
model/models/replicate_api_models.py +7 -148
serve/Ksort.py +29 -0
serve/leaderboard.py +4 -0
serve/update_skill.py +11 -3
serve/update_skill_video.py +11 -3
serve/upload.py +13 -5

app.py CHANGED Viewed

@@ -7,12 +7,8 @@ from model.model_manager import ModelManager
 from pathlib import Path
 from serve.constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR
 def build_combine_demo(models, elo_results_file, leaderboard_table_file):
-        # gr.themes.Default(),
-        # gr.themes.Soft(),
-        # gr.Theme.from_hub('gary109/HaleyCH_Theme'),
-        # gr.Theme.from_hub('EveryPizza/Cartoony-Gradio-Theme')
-        # gr.themes.Default(primary_hue="red", secondary_hue="pink")
     with gr.Blocks(
         title="Play with Open Vision Models",
         theme=gr.themes.Default(),
@@ -22,21 +18,21 @@ def build_combine_demo(models, elo_results_file, leaderboard_table_file):
             with gr.Tab("Image Generation", id=0):
                 with gr.Tabs() as tabs_ig:
                     with gr.Tab("Generation Leaderboard", id=0):
-                        # build_leaderboard_tab(elo_results_file['t2i_generation'], leaderboard_table_file['t2i_generation'])
                         build_leaderboard_tab()
                     with gr.Tab("Generation Arena (battle)", id=1):
                         build_side_by_side_ui_anony(models)
             with gr.Tab("Video Generation", id=1):
                 with gr.Tabs() as tabs_ig:
                     with gr.Tab("Generation Leaderboard", id=0):
-                        # build_leaderboard_tab(elo_results_file['t2i_generation'], leaderboard_table_file['t2i_generation'])
                         build_leaderboard_video_tab()
                     with gr.Tab("Generation Arena (battle)", id=1):
                         build_side_by_side_video_ui_anony(models)
             with gr.Tab("Contributor", id=2):
                 build_leaderboard_contributor()
     return demo
@@ -44,6 +40,7 @@ def load_elo_results(elo_results_dir):
     from collections import defaultdict
     elo_results_file = defaultdict(lambda: None)
     leaderboard_table_file = defaultdict(lambda: None)
     if elo_results_dir is not None:
         elo_results_dir = Path(elo_results_dir)
         elo_results_file = {}

 from pathlib import Path
 from serve.constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR
 def build_combine_demo(models, elo_results_file, leaderboard_table_file):
     with gr.Blocks(
         title="Play with Open Vision Models",
         theme=gr.themes.Default(),
             with gr.Tab("Image Generation", id=0):
                 with gr.Tabs() as tabs_ig:
                     with gr.Tab("Generation Leaderboard", id=0):
                         build_leaderboard_tab()
                     with gr.Tab("Generation Arena (battle)", id=1):
                         build_side_by_side_ui_anony(models)
             with gr.Tab("Video Generation", id=1):
                 with gr.Tabs() as tabs_ig:
                     with gr.Tab("Generation Leaderboard", id=0):
                         build_leaderboard_video_tab()
                     with gr.Tab("Generation Arena (battle)", id=1):
                         build_side_by_side_video_ui_anony(models)
             with gr.Tab("Contributor", id=2):
                 build_leaderboard_contributor()
     return demo
     from collections import defaultdict
     elo_results_file = defaultdict(lambda: None)
     leaderboard_table_file = defaultdict(lambda: None)
     if elo_results_dir is not None:
         elo_results_dir = Path(elo_results_dir)
         elo_results_file = {}

model/matchmaker.py CHANGED Viewed

@@ -24,35 +24,41 @@ def create_ssh_matchmaker_client(server, port, user, password):
     transport.set_keepalive(60)
     sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
 def is_connected():
     global ssh_matchmaker_client, sftp_matchmaker_client
     if ssh_matchmaker_client is None or sftp_matchmaker_client is None:
         return False
-    # 检查SSH连接是否正常
     if not ssh_matchmaker_client.get_transport().is_active():
         return False
-    # 检查SFTP连接是否正常
     try:
-        sftp_matchmaker_client.listdir('.')  # 尝试列出根目录
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_matchmaker_client
     if not is_connected():
@@ -66,6 +72,7 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     with sftp_matchmaker_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_matchmaker_client
     if not is_connected():
@@ -107,7 +114,7 @@ def matchmaker(num_players, k_group=4, not_run=[]):
     ucb_scores = ucb_score(trueskill_diff, total_comparisons, n)
     # Exclude self, select opponent with highest UCB score
-    ucb_scores[selected_player] = -float('inf')  # minimize the score for the selected player to exclude it
     ucb_scores[not_run] = -float('inf')
     opponents = np.argsort(ucb_scores)[-k_group + 1:].tolist()
@@ -117,4 +124,3 @@ def matchmaker(num_players, k_group=4, not_run=[]):
     random.shuffle(model_ids)
     return model_ids

     transport.set_keepalive(60)
     sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
 def is_connected():
     global ssh_matchmaker_client, sftp_matchmaker_client
     if ssh_matchmaker_client is None or sftp_matchmaker_client is None:
         return False
     if not ssh_matchmaker_client.get_transport().is_active():
         return False
     try:
+        sftp_matchmaker_client.listdir('.')
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_matchmaker_client
     if not is_connected():
     with sftp_matchmaker_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_matchmaker_client
     if not is_connected():
     ucb_scores = ucb_score(trueskill_diff, total_comparisons, n)
     # Exclude self, select opponent with highest UCB score
+    ucb_scores[selected_player] = -float('inf')
     ucb_scores[not_run] = -float('inf')
     opponents = np.argsort(ucb_scores)[-k_group + 1:].tolist()
     random.shuffle(model_ids)
     return model_ids

model/matchmaker_video.py CHANGED Viewed

@@ -13,6 +13,7 @@ trueskill_env = TrueSkill()
 ssh_matchmaker_client = None
 sftp_matchmaker_client = None
 def create_ssh_matchmaker_client(server, port, user, password):
     global ssh_matchmaker_client, sftp_matchmaker_client
     ssh_matchmaker_client = paramiko.SSHClient()
@@ -24,35 +25,41 @@ def create_ssh_matchmaker_client(server, port, user, password):
     transport.set_keepalive(60)
     sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
 def is_connected():
     global ssh_matchmaker_client, sftp_matchmaker_client
     if ssh_matchmaker_client is None or sftp_matchmaker_client is None:
         return False
-    # 检查SSH连接是否正常
     if not ssh_matchmaker_client.get_transport().is_active():
         return False
-    # 检查SFTP连接是否正常
     try:
-        sftp_matchmaker_client.listdir('.')  # 尝试列出根目录
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_matchmaker_client
     if not is_connected():
@@ -66,6 +73,7 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     with sftp_matchmaker_client.open(SSH_VIDEO_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_matchmaker_client
     if not is_connected():
@@ -95,7 +103,7 @@ def matchmaker_video(num_players, k_group=4):
     ucb_scores = ucb_score(trueskill_diff, total_comparisons, n)
     # Exclude self, select opponent with highest UCB score
-    ucb_scores[selected_player] = -float('inf')  # minimize the score for the selected player to exclude it
     excluded_players_1 = [num_players-1, num_players-4]
     excluded_players_2 = [num_players-2, num_players-3, num_players-5]
@@ -126,4 +134,3 @@ def matchmaker_video(num_players, k_group=4):
     random.shuffle(model_ids)
     return model_ids

 ssh_matchmaker_client = None
 sftp_matchmaker_client = None
 def create_ssh_matchmaker_client(server, port, user, password):
     global ssh_matchmaker_client, sftp_matchmaker_client
     ssh_matchmaker_client = paramiko.SSHClient()
     transport.set_keepalive(60)
     sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
 def is_connected():
     global ssh_matchmaker_client, sftp_matchmaker_client
     if ssh_matchmaker_client is None or sftp_matchmaker_client is None:
         return False
     if not ssh_matchmaker_client.get_transport().is_active():
         return False
     try:
+        sftp_matchmaker_client.listdir('.')
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_matchmaker_client
     if not is_connected():
     with sftp_matchmaker_client.open(SSH_VIDEO_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_matchmaker_client
     if not is_connected():
     ucb_scores = ucb_score(trueskill_diff, total_comparisons, n)
     # Exclude self, select opponent with highest UCB score
+    ucb_scores[selected_player] = -float('inf')
     excluded_players_1 = [num_players-1, num_players-4]
     excluded_players_2 = [num_players-2, num_players-3, num_players-5]
     random.shuffle(model_ids)
     return model_ids

model/model_manager.py CHANGED Viewed

@@ -58,10 +58,8 @@ class ModelManager:
     def generate_image_ig_api(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
         return result
     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 == "":
             from .matchmaker import matchmaker
@@ -73,13 +71,11 @@ class ModelManager:
         else:
             model_names = [model_A, model_B, model_C, model_D]
         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]
@@ -156,7 +152,6 @@ class ModelManager:
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3], prompt
     def generate_image_ig_parallel(self, prompt, model_A, model_B):
         model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -165,14 +160,12 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
     @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_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:
@@ -182,7 +175,6 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
     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)

     def generate_image_ig_api(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
         return result
     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 == "":
             from .matchmaker import matchmaker
         else:
             model_names = [model_A, model_B, model_C, model_D]
         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]
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3], prompt
     def generate_image_ig_parallel(self, prompt, model_A, model_B):
         model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
             results = [future.result() for future in futures]
         return results[0], results[1]
     @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_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:
             results = [future.result() for future in futures]
         return results[0], results[1]
     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)

model/model_registry.py CHANGED Viewed

@@ -68,188 +68,3 @@ def get_video_model_description_md(model_list):
             model_description_md += "\n"
         ct += 1
     return model_description_md
-register_model_info(
-    ["imagenhub_LCM_generation", "fal_LCM_text2image"],
-    "LCM",
-    "https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7",
-    "Latent Consistency Models.",
-)
-register_model_info(
-    ["fal_LCM(v1.5/XL)_text2image"],
-    "LCM(v1.5/XL)",
-    "https://fal.ai/models/fast-lcm-diffusion-turbo",
-    "Latent Consistency Models (v1.5/XL)",
-)
-register_model_info(
-    ["imagenhub_PlayGroundV2_generation", 'playground_PlayGroundV2_generation'],
-    "Playground v2",
-    "https://huggingface.co/playgroundai/playground-v2-1024px-aesthetic",
-    "Playground v2 – 1024px Aesthetic Model",
-)
-register_model_info(
-    ["imagenhub_PlayGroundV2.5_generation", 'playground_PlayGroundV2.5_generation'],
-    "Playground v2.5",
-    "https://huggingface.co/playgroundai/playground-v2.5-1024px-aesthetic",
-    "Playground v2.5 is the state-of-the-art open-source model in aesthetic quality",
-)
-register_model_info(
-    ["imagenhub_OpenJourney_generation"],
-    "Openjourney",
-    "https://huggingface.co/prompthero/openjourney",
-    "Openjourney is an open source Stable Diffusion fine tuned model on Midjourney images, by PromptHero.",
-)
-register_model_info(
-    ["imagenhub_SDXLTurbo_generation", "fal_SDXLTurbo_text2image"],
-    "SDXLTurbo",
-    "https://huggingface.co/stabilityai/sdxl-turbo",
-    "SDXL-Turbo is a fast generative text-to-image model.",
-)
-register_model_info(
-    ["imagenhub_SDXL_generation", "fal_SDXL_text2image"],
-    "SDXL",
-    "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0",
-    "SDXL is a Latent Diffusion Model that uses two fixed, pretrained text encoders.",
-)
-register_model_info(
-    ["imagenhub_PixArtAlpha_generation"],
-    "PixArtAlpha",
-    "https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS",
-    "Pixart-α consists of pure transformer blocks for latent diffusion.",
-)
-register_model_info(
-    ["imagenhub_PixArtSigma_generation", "fal_PixArtSigma_text2image"],
-    "PixArtSigma",
-    "https://github.com/PixArt-alpha/PixArt-sigma",
-    "Improved version of Pixart-α.",
-)
-register_model_info(
-    ["imagenhub_SDXLLightning_generation", "fal_SDXLLightning_text2image"],
-    "SDXL-Lightning",
-    "https://huggingface.co/ByteDance/SDXL-Lightning",
-    "SDXL-Lightning is a lightning-fast text-to-image generation model.",
-)
-register_model_info(
-    ["imagenhub_StableCascade_generation", "fal_StableCascade_text2image"],
-    "StableCascade",
-    "https://huggingface.co/stabilityai/stable-cascade",
-    "StableCascade is built upon the Würstchen architecture and working at a much smaller latent space.",
-)
-# regist image edition models
-register_model_info(
-    ["imagenhub_CycleDiffusion_edition"],
-    "CycleDiffusion",
-    "https://github.com/ChenWu98/cycle-diffusion?tab=readme-ov-file",
-    "A latent space for stochastic diffusion models.",
-)
-register_model_info(
-    ["imagenhub_Pix2PixZero_edition"],
-    "Pix2PixZero",
-    "https://pix2pixzero.github.io/",
-    "A zero-shot Image-to-Image translation model.",
-)
-register_model_info(
-    ["imagenhub_Prompt2prompt_edition"],
-    "Prompt2prompt",
-    "https://prompt-to-prompt.github.io/",
-    "Image Editing with Cross-Attention Control.",
-)
-register_model_info(
-    ["imagenhub_InstructPix2Pix_edition"],
-    "InstructPix2Pix",
-    "https://www.timothybrooks.com/instruct-pix2pix",
-    "An instruction-based image editing model.",
-)
-register_model_info(
-    ["imagenhub_MagicBrush_edition"],
-    "MagicBrush",
-    "https://osu-nlp-group.github.io/MagicBrush/",
-    "Manually Annotated Dataset for Instruction-Guided Image Editing.",
-)
-register_model_info(
-    ["imagenhub_PNP_edition"],
-    "PNP",
-    "https://github.com/MichalGeyer/plug-and-play",
-    "Plug-and-Play Diffusion Features for Text-Driven Image-to-Image Translation.",
-)
-register_model_info(
-    ["imagenhub_InfEdit_edition"],
-    "InfEdit",
-    "https://sled-group.github.io/InfEdit/",
-    "Inversion-Free Image Editing with Natural Language.",
-)
-register_model_info(
-    ["imagenhub_CosXLEdit_edition"],
-    "CosXLEdit",
-    "https://huggingface.co/stabilityai/cosxl",
-    "An instruction-based image editing model from SDXL.",
-)
-register_model_info(
-    ["fal_stable-cascade_text2image"],
-    "StableCascade",
-    "https://fal.ai/models/stable-cascade/api",
-    "StableCascade is a generative model that can generate high-quality images from text prompts.",
-)
-register_model_info(
-    ["fal_AnimateDiff_text2video"],
-    "AnimateDiff",
-    "https://fal.ai/models/fast-animatediff-t2v",
-    "AnimateDiff is a text-driven models that produce diverse and personalized animated images.",
-)
-register_model_info(
-    ["fal_AnimateDiffTurbo_text2video"],
-    "AnimateDiff Turbo",
-    "https://fal.ai/models/fast-animatediff-t2v-turbo",
-    "AnimateDiff Turbo is a lightning version of AnimateDiff.",
-)
-register_model_info(
-    ["videogenhub_LaVie_generation"],
-    "LaVie",
-    "https://github.com/Vchitect/LaVie",
-    "LaVie is a video generation model with cascaded latent diffusion models.",
-)
-register_model_info(
-    ["videogenhub_VideoCrafter2_generation"],
-    "VideoCrafter2",
-    "https://ailab-cvc.github.io/videocrafter2/",
-    "VideoCrafter2 is a T2V model that disentangling motion from appearance.",
-)
-register_model_info(
-    ["videogenhub_ModelScope_generation"],
-    "ModelScope",
-    "https://arxiv.org/abs/2308.06571",
-    "ModelScope is a a T2V synthesis model that evolves from a T2I synthesis model.",
-)
-register_model_info(
-    ["videogenhub_OpenSora_generation"],
-    "OpenSora",
-    "https://github.com/hpcaitech/Open-Sora",
-    "A community-driven opensource implementation of Sora.",
-)

             model_description_md += "\n"
         ct += 1
     return model_description_md

model/models/__init__.py CHANGED Viewed

@@ -37,15 +37,6 @@ IMAGE_GENERATION_MODELS = [
                             "replicate_FLUX.1-dev_text2image",
                             ]
-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']
 VIDEO_GENERATION_MODELS = ['replicate_Zeroscope-v2-xl_text2video',
                             'replicate_Animate-Diff_text2video',
                             'replicate_OpenSora_text2video',
@@ -59,22 +50,15 @@ VIDEO_GENERATION_MODELS = ['replicate_Zeroscope-v2-xl_text2video',
                             'other_Sora_text2video',
                            ]
 def load_pipeline(model_name):
     """
     Load a model pipeline based on the model name
     Args:
         model_name (str): The name of the model to load, should be of the form {source}_{name}_{type}
-        the source can be either imagenhub or playground
-        the name is the name of the model used to load the model
-        the type is the type of the model, either generation or edition
     """
     model_source, model_name, model_type = model_name.split("_")
-    # if model_source == "imagenhub":
-    #     pipe = load_imagenhub_model(model_name, model_type)
-    # elif model_source == "fal":
-    #     pipe = load_fal_model(model_name, model_type)
-    # elif model_source == "videogenhub":
-    #     pipe = load_videogenhub_model(model_name)
     if model_source == "replicate":
         pipe = load_replicate_model(model_name, model_type)
     elif model_source == "huggingface":

                             "replicate_FLUX.1-dev_text2image",
                             ]
 VIDEO_GENERATION_MODELS = ['replicate_Zeroscope-v2-xl_text2video',
                             'replicate_Animate-Diff_text2video',
                             'replicate_OpenSora_text2video',
                             'other_Sora_text2video',
                            ]
 def load_pipeline(model_name):
     """
     Load a model pipeline based on the model name
     Args:
         model_name (str): The name of the model to load, should be of the form {source}_{name}_{type}
     """
     model_source, model_name, model_type = model_name.split("_")
     if model_source == "replicate":
         pipe = load_replicate_model(model_name, model_type)
     elif model_source == "huggingface":

model/models/openai_api_models.py CHANGED Viewed

@@ -12,7 +12,6 @@ class OpenaiModel():
         self.model_type = model_type
     def __call__(self, *args, **kwargs):
         if self.model_type == "text2image":
             assert "prompt" in kwargs, "prompt is required for text2image model"
@@ -47,7 +46,6 @@ class OpenaiModel():
             raise ValueError("model_type must be text2image or image2image")
 def load_openai_model(model_name, model_type):
     return OpenaiModel(model_name, model_type)

         self.model_type = model_type
     def __call__(self, *args, **kwargs):
         if self.model_type == "text2image":
             assert "prompt" in kwargs, "prompt is required for text2image model"
             raise ValueError("model_type must be text2image or image2image")
 def load_openai_model(model_name, model_type):
     return OpenaiModel(model_name, model_type)

model/models/other_api_models.py CHANGED Viewed

@@ -4,6 +4,7 @@ import os
 from PIL import Image
 import io, time
 class OtherModel():
     def __init__(self, model_name, model_type):
         self.model_name = model_name
@@ -75,6 +76,8 @@ class OtherModel():
         else:
             raise ValueError("model_type must be text2image")
 def load_other_model(model_name, model_type):
     return OtherModel(model_name, model_type)
@@ -86,30 +89,3 @@ if __name__ == "__main__":
     result = pipe(prompt="An Impressionist illustration depicts a river winding through a meadow ")
     print(result)
     exit()
-    # key = os.environ.get('MIDJOURNEY_KEY')
-    # prompt = "a good girl"
-    # conn = http.client.HTTPSConnection("xdai.online")
-    # payload = json.dumps({
-    # "messages": [
-    #     {
-    #         "role": "user",
-    #         "content": "{}".format(prompt)
-    #     }
-    # ],
-    # "stream": True,
-    # "model": "luma-video",
-    # # "model": "pika-text-to-video",
-    # })
-    # headers = {
-    #     'Authorization': "Bearer {}".format(key),
-    #     'Content-Type': 'application/json'
-    # }
-    # conn.request("POST", "/v1/chat/completions", payload, headers)
-    # res = conn.getresponse()
-    # data = res.read()
-    # info = data.decode("utf-8")
-    # print(data.decode("utf-8"))

 from PIL import Image
 import io, time
 class OtherModel():
     def __init__(self, model_name, model_type):
         self.model_name = model_name
         else:
             raise ValueError("model_type must be text2image")
 def load_other_model(model_name, model_type):
     return OtherModel(model_name, model_type)
     result = pipe(prompt="An Impressionist illustration depicts a river winding through a meadow ")
     print(result)
     exit()

model/models/replicate_api_models.py CHANGED Viewed

@@ -40,11 +40,11 @@ Replicate_MODEl_NAME_MAP = {
     "FLUX.1-dev": "black-forest-labs/flux-dev",
     }
 class ReplicateModel():
     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):
         if self.model_type == "text2image":
@@ -179,155 +179,14 @@ class ReplicateModel():
         else:
             raise ValueError("model_type must be text2image or image2image")
 def load_replicate_model(model_name, model_type):
     return ReplicateModel(model_name, model_type)
 if __name__ == "__main__":
-    import replicate
-    import time
-    import concurrent.futures
-    import os, shutil, re
-    import requests
-    from moviepy.editor import VideoFileClip
-    # model_name = 'replicate_zeroscope-v2-xl_text2video'
-    # model_name = 'replicate_Damo-Text-to-Video_text2video'
-    # model_name = 'replicate_Animate-Diff_text2video'
-    # model_name = 'replicate_open-sora_text2video'
-    # model_name = 'replicate_lavie_text2video'
-    # model_name = 'replicate_video-crafter_text2video'
-    # model_name = 'replicate_stable-video-diffusion_text2video'
-    # model_source, model_name, model_type = model_name.split("_")
-    # pipe = load_replicate_model(model_name, model_type)
-    # prompt = "Clown fish swimming in a coral reef, beautiful, 8k, perfect, award winning, national geographic"
-    # result = pipe(prompt=prompt)
-    # # 文件复制
-    source_folder = '/mnt/data/lizhikai/ksort_video_cache/Pika-v1.0add/'
-    destination_folder = '/mnt/data/lizhikai/ksort_video_cache/Advance/'
-    special_char = 'output'
-    for dirpath, dirnames, filenames in os.walk(source_folder):
-        for dirname in dirnames:
-            des_dirname = "output-"+dirname[-3:]
-            print(des_dirname)
-            if special_char in dirname:
-                model_name = ["Pika-v1.0"]
-                for name in model_name:
-                    source_file_path = os.path.join(source_folder, os.path.join(dirname, name+".mp4"))
-                    print(source_file_path)
-                    destination_file_path = os.path.join(destination_folder, os.path.join(des_dirname, name+".mp4"))
-                    print(destination_file_path)
-                    shutil.copy(source_file_path, destination_file_path)
-    # 视频裁剪
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Runway-Gen3/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Runway-Gen2/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Pika-Beta/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Pika-v1/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Sora/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Pika-v1.0add/'
-    # special_char = 'output'
-    # num = 0
-    # for dirpath, dirnames, filenames in os.walk(root_dir):
-    #     for dirname in dirnames:
-    #         # 如果文件夹名称中包含指定的特殊字符
-    #         if special_char in dirname:
-    #             num = num+1
-    #             print(num)
-    #             if num < 0:
-    #                 continue
-    #             video_path = os.path.join(root_dir, (os.path.join(dirname, f"{dirname}.mp4")))
-    #             out_video_path = os.path.join(root_dir, (os.path.join(dirname, f"Pika-v1.0.mp4")))
-    #             print(video_path)
-    #             print(out_video_path)
-    #             video = VideoFileClip(video_path)
-    #             width, height = video.size
-    #             center_x, center_y = width // 2, height // 2
-    #             new_width, new_height = 512, 512
-    #             cropped_video = video.crop(x_center=center_x, y_center=center_y, width=min(width, height), height=min(width, height))
-    #             resized_video = cropped_video.resize(newsize=(new_width, new_height))
-    #             resized_video.write_videofile(out_video_path, codec='libx264', fps=video.fps)
-    #             os.remove(video_path)
-    # file_path = '/home/lizhikai/webvid_prompt100.txt'
-    # str_list = []
-    # with open(file_path, 'r', encoding='utf-8') as file:
-    #     for line in file:
-    #         str_list.append(line.strip())
-    #         if len(str_list) == 100:
-    #             break
-    # 生成代码
-    # def generate_image_ig_api(prompt, model_name):
-    #     model_source, model_name, model_type = model_name.split("_")
-    #     pipe = load_replicate_model(model_name, model_type)
-    #     result = pipe(prompt=prompt)
-    #     return result
-    # model_names = ['replicate_Zeroscope-v2-xl_text2video',
-    #                     #    'replicate_Damo-Text-to-Video_text2video',
-    #                        'replicate_Animate-Diff_text2video',
-    #                        'replicate_OpenSora_text2video',
-    #                        'replicate_LaVie_text2video',
-    #                        'replicate_VideoCrafter2_text2video',
-    #                        'replicate_Stable-Video-Diffusion_text2video',
-    #                        ]
-    # save_names = []
-    # for name in model_names:
-    #     model_source, model_name, model_type = name.split("_")
-    #     save_names.append(model_name)
-    # # 遍历根目录及其子目录
-    # # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Runway-Gen3/'
-    # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Runway-Gen2/'
-    # # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Pika-Beta/'
-    # # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Pika-v1/'
-    # # root_dir = '/mnt/data/lizhikai/ksort_video_cache/Sora/'
-    # special_char = 'output'
-    # num = 0
-    # for dirpath, dirnames, filenames in os.walk(root_dir):
-    #     for dirname in dirnames:
-    #         # 如果文件夹名称中包含指定的特殊字符
-    #         if special_char in dirname:
-    #             num = num+1
-    #             print(num)
-    #             if num < 0:
-    #                 continue
-    #             str_list = []
-    #             prompt_path = os.path.join(root_dir, (os.path.join(dirname, "prompt.txt")))
-    #             print(prompt_path)
-    #             with open(prompt_path, 'r', encoding='utf-8') as file:
-    #                 for line in file:
-    #                     str_list.append(line.strip())
-    #             prompt = str_list[0]
-    #             print(prompt)
-    #             with concurrent.futures.ThreadPoolExecutor() as executor:
-    #                 futures = [executor.submit(generate_image_ig_api, prompt, model) for model in model_names]
-    #                 results = [future.result() for future in futures]
-    #             # 下载视频并保存
-    #             repeat_num = 5
-    #             for j, url in enumerate(results):
-    #                 while 1:
-    #                     time.sleep(1)
-    #                     response = requests.get(url, stream=True)
-    #                     if response.status_code == 200:
-    #                         file_path = os.path.join(os.path.join(root_dir, dirname), f'{save_names[j]}.mp4')
-    #                         with open(file_path, 'wb') as file:
-    #                             for chunk in response.iter_content(chunk_size=8192):
-    #                                 file.write(chunk)
-    #                         print(f"视频 {j} 已保存到 {file_path}")
-    #                         break
-    #                     else:
-    #                         repeat_num = repeat_num - 1
-    #                         if repeat_num == 0:
-    #                             print(f"视频 {j} 保存失败")
-    #                             # raise ValueError("Video request failed.")
-    #                         continue

     "FLUX.1-dev": "black-forest-labs/flux-dev",
     }
 class ReplicateModel():
     def __init__(self, model_name, model_type):
         self.model_name = model_name
         self.model_type = model_type
     def __call__(self, *args, **kwargs):
         if self.model_type == "text2image":
         else:
             raise ValueError("model_type must be text2image or image2image")
 def load_replicate_model(model_name, model_type):
     return ReplicateModel(model_name, model_type)
 if __name__ == "__main__":
+    model_name = 'replicate_zeroscope-v2-xl_text2video'
+    model_source, model_name, model_type = model_name.split("_")
+    pipe = load_replicate_model(model_name, model_type)
+    prompt = "Clown fish swimming in a coral reef, beautiful, 8k, perfect, award winning, national geographic"
+    result = pipe(prompt=prompt)

serve/Ksort.py CHANGED Viewed

@@ -8,6 +8,7 @@ from .utils import disable_btn, enable_btn, invisible_btn
 from .upload import create_remote_directory, upload_ssh_all, upload_ssh_data
 import json
 def reset_level(Top_btn):
     if Top_btn == "Top 1":
         level = 0
@@ -19,6 +20,7 @@ def reset_level(Top_btn):
         level = 3
     return level
 def reset_rank(windows, rank, vote_level):
     if windows == "Model A":
         rank[0] = vote_level
@@ -30,6 +32,7 @@ def reset_rank(windows, rank, vote_level):
         rank[3] = vote_level
     return rank
 def reset_btn_rank(windows, rank, btn, vote_level):
     if windows == "Model A" and btn == "1":
         rank[0] = 0
@@ -73,6 +76,7 @@ def reset_btn_rank(windows, rank, btn, vote_level):
         vote_level = 3
     return (rank, vote_level)
 def reset_vote_text(rank):
     rank_str = ""
     for i in range(len(rank)):
@@ -83,24 +87,28 @@ def reset_vote_text(rank):
         rank_str = rank_str + "  "
     return rank_str
 def clear_rank(rank, vote_level):
     for i in range(len(rank)):
         rank[i] = None
     vote_level = 0
     return rank, vote_level
 def revote_windows(generate_ig0, generate_ig1, generate_ig2, generate_ig3, rank, vote_level):
     for i in range(len(rank)):
         rank[i] = None
     vote_level = 0
     return generate_ig0, generate_ig1, generate_ig2, generate_ig3, rank, vote_level
 def reset_submit(rank):
     for i in range(len(rank)):
         if rank[i] == None:
             return disable_btn
     return enable_btn
 def reset_mode(mode):
     if mode == "Best":
@@ -116,8 +124,12 @@ def reset_mode(mode):
                 (gr.Textbox(value="Best", visible=False, interactive=False),)
     else:
         raise ValueError("Undefined mode")
 def reset_chatbot(mode, generate_ig0, generate_ig1, generate_ig2, generate_ig3):
     return generate_ig0, generate_ig1, generate_ig2, generate_ig3
 def get_json_filename(conv_id):
     output_dir = f'{KSORT_IMAGE_DIR}/{conv_id}/json/'
     if not os.path.exists(output_dir):
@@ -127,6 +139,7 @@ def get_json_filename(conv_id):
     print(output_file)
     return output_file
 def get_img_filename(conv_id, i):
     output_dir = f'{KSORT_IMAGE_DIR}/{conv_id}/image/'
     if not os.path.exists(output_dir):
@@ -135,6 +148,7 @@ def get_img_filename(conv_id, i):
     print(output_file)
     return output_file
 def vote_submit(states, textbox, rank, request: gr.Request):
     conv_id = states[0].conv_id
@@ -149,6 +163,7 @@ def vote_submit(states, textbox, rank, request: gr.Request):
         }
         fout.write(json.dumps(data) + "\n")
 def vote_ssh_submit(states, textbox, rank, user_name, user_institution):
     conv_id = states[0].conv_id
     output_dir = create_remote_directory(conv_id)
@@ -167,6 +182,7 @@ def vote_ssh_submit(states, textbox, rank, user_name, user_institution):
     from .update_skill import update_skill
     update_skill(rank, [x.model_name for x in states])
 def vote_video_ssh_submit(states, textbox, prompt_path, rank, user_name, user_institution):
     conv_id = states[0].conv_id
     output_dir = create_remote_directory(conv_id, video=True)
@@ -186,6 +202,7 @@ def vote_video_ssh_submit(states, textbox, prompt_path, rank, user_name, user_in
     from .update_skill_video import update_skill_video
     update_skill_video(rank, [x.model_name for x in states])
 def submit_response_igm(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, rank, user_name, user_institution, request: gr.Request
 ):
@@ -205,6 +222,8 @@ def submit_response_igm(
             gr.Markdown(state2.model_name, visible=True),
             gr.Markdown(state3.model_name, visible=True)
             ) + (disable_btn,)
 def submit_response_vg(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, prompt_path, rank, user_name, user_institution, request: gr.Request
 ):
@@ -223,6 +242,8 @@ def submit_response_vg(
             gr.Markdown(state2.model_name, visible=True),
             gr.Markdown(state3.model_name, visible=True)
             ) + (disable_btn,)
 def submit_response_rank_igm(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, rank, right_vote_text, user_name, user_institution, request: gr.Request
 ):
@@ -246,6 +267,8 @@ def submit_response_rank_igm(
                 )
     else:
         return (enable_btn,) * 16 + (enable_btn,) * 3 + ("wrong",) + (gr.Markdown("", visible=False),) * 4
 def submit_response_rank_vg(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, prompt_path, rank, right_vote_text, user_name, user_institution, request: gr.Request
 ):
@@ -269,6 +292,7 @@ def submit_response_rank_vg(
     else:
         return (enable_btn,) * 16 + (enable_btn,) * 3 + ("wrong",) + (gr.Markdown("", visible=False),) * 4
 def text_response_rank_igm(generate_ig0, generate_ig1, generate_ig2, generate_ig3, Top1_text, Top2_text, Top3_text, Top4_text, vote_textbox):
     rank_list = [char for char in vote_textbox if char.isdigit()]
     generate_ig = [generate_ig0, generate_ig1, generate_ig2, generate_ig3]
@@ -318,6 +342,7 @@ def text_response_rank_igm(generate_ig0, generate_ig1, generate_ig2, generate_ig
     return chatbot + [rank_str] + ["right"] + [rank]
 def text_response_rank_vg(vote_textbox):
     rank_list = [char for char in vote_textbox if char.isdigit()]
     rank = [None, None, None, None]
@@ -336,6 +361,7 @@ def text_response_rank_vg(vote_textbox):
     return [rank_str] + ["right"] + [rank]
 def add_foreground(image, vote_level, Top1_text, Top2_text, Top3_text, Top4_text):
     base_image = Image.fromarray(image).convert("RGBA")
     base_image = base_image.resize((512, 512), Image.ANTIALIAS)
@@ -369,12 +395,15 @@ def add_foreground(image, vote_level, Top1_text, Top2_text, Top3_text, Top4_text
     base_image = base_image.convert("RGB")
     return base_image
 def add_green_border(image):
     border_color = (0, 255, 0)  # RGB for green
     border_size = 10  # Size of the border
     img_with_border = ImageOps.expand(image, border=border_size, fill=border_color)
     return img_with_border
 def check_textbox(textbox):
     if textbox=="":
         return False

 from .upload import create_remote_directory, upload_ssh_all, upload_ssh_data
 import json
 def reset_level(Top_btn):
     if Top_btn == "Top 1":
         level = 0
         level = 3
     return level
 def reset_rank(windows, rank, vote_level):
     if windows == "Model A":
         rank[0] = vote_level
         rank[3] = vote_level
     return rank
 def reset_btn_rank(windows, rank, btn, vote_level):
     if windows == "Model A" and btn == "1":
         rank[0] = 0
         vote_level = 3
     return (rank, vote_level)
 def reset_vote_text(rank):
     rank_str = ""
     for i in range(len(rank)):
         rank_str = rank_str + "  "
     return rank_str
 def clear_rank(rank, vote_level):
     for i in range(len(rank)):
         rank[i] = None
     vote_level = 0
     return rank, vote_level
 def revote_windows(generate_ig0, generate_ig1, generate_ig2, generate_ig3, rank, vote_level):
     for i in range(len(rank)):
         rank[i] = None
     vote_level = 0
     return generate_ig0, generate_ig1, generate_ig2, generate_ig3, rank, vote_level
 def reset_submit(rank):
     for i in range(len(rank)):
         if rank[i] == None:
             return disable_btn
     return enable_btn
 def reset_mode(mode):
     if mode == "Best":
                 (gr.Textbox(value="Best", visible=False, interactive=False),)
     else:
         raise ValueError("Undefined mode")
 def reset_chatbot(mode, generate_ig0, generate_ig1, generate_ig2, generate_ig3):
     return generate_ig0, generate_ig1, generate_ig2, generate_ig3
 def get_json_filename(conv_id):
     output_dir = f'{KSORT_IMAGE_DIR}/{conv_id}/json/'
     if not os.path.exists(output_dir):
     print(output_file)
     return output_file
 def get_img_filename(conv_id, i):
     output_dir = f'{KSORT_IMAGE_DIR}/{conv_id}/image/'
     if not os.path.exists(output_dir):
     print(output_file)
     return output_file
 def vote_submit(states, textbox, rank, request: gr.Request):
     conv_id = states[0].conv_id
         }
         fout.write(json.dumps(data) + "\n")
 def vote_ssh_submit(states, textbox, rank, user_name, user_institution):
     conv_id = states[0].conv_id
     output_dir = create_remote_directory(conv_id)
     from .update_skill import update_skill
     update_skill(rank, [x.model_name for x in states])
 def vote_video_ssh_submit(states, textbox, prompt_path, rank, user_name, user_institution):
     conv_id = states[0].conv_id
     output_dir = create_remote_directory(conv_id, video=True)
     from .update_skill_video import update_skill_video
     update_skill_video(rank, [x.model_name for x in states])
 def submit_response_igm(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, rank, user_name, user_institution, request: gr.Request
 ):
             gr.Markdown(state2.model_name, visible=True),
             gr.Markdown(state3.model_name, visible=True)
             ) + (disable_btn,)
 def submit_response_vg(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, prompt_path, rank, user_name, user_institution, request: gr.Request
 ):
             gr.Markdown(state2.model_name, visible=True),
             gr.Markdown(state3.model_name, visible=True)
             ) + (disable_btn,)
 def submit_response_rank_igm(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, rank, right_vote_text, user_name, user_institution, request: gr.Request
 ):
                 )
     else:
         return (enable_btn,) * 16 + (enable_btn,) * 3 + ("wrong",) + (gr.Markdown("", visible=False),) * 4
 def submit_response_rank_vg(
     state0, state1, state2, state3, model_selector0, model_selector1, model_selector2, model_selector3, textbox, prompt_path, rank, right_vote_text, user_name, user_institution, request: gr.Request
 ):
     else:
         return (enable_btn,) * 16 + (enable_btn,) * 3 + ("wrong",) + (gr.Markdown("", visible=False),) * 4
 def text_response_rank_igm(generate_ig0, generate_ig1, generate_ig2, generate_ig3, Top1_text, Top2_text, Top3_text, Top4_text, vote_textbox):
     rank_list = [char for char in vote_textbox if char.isdigit()]
     generate_ig = [generate_ig0, generate_ig1, generate_ig2, generate_ig3]
     return chatbot + [rank_str] + ["right"] + [rank]
 def text_response_rank_vg(vote_textbox):
     rank_list = [char for char in vote_textbox if char.isdigit()]
     rank = [None, None, None, None]
     return [rank_str] + ["right"] + [rank]
 def add_foreground(image, vote_level, Top1_text, Top2_text, Top3_text, Top4_text):
     base_image = Image.fromarray(image).convert("RGBA")
     base_image = base_image.resize((512, 512), Image.ANTIALIAS)
     base_image = base_image.convert("RGB")
     return base_image
 def add_green_border(image):
     border_color = (0, 255, 0)  # RGB for green
     border_size = 10  # Size of the border
     img_with_border = ImageOps.expand(image, border=border_size, fill=border_color)
     return img_with_border
 def check_textbox(textbox):
     if textbox=="":
         return False

serve/leaderboard.py CHANGED Viewed

@@ -40,12 +40,14 @@ def make_leaderboard_md():
 """
     return leaderboard_md
 def make_leaderboard_video_md():
     leaderboard_md = f"""
 # 🏆 K-Sort Arena Leaderboard (Text-to-Video Generation)
 """
     return leaderboard_md
 def model_hyperlink(model_name, link):
     return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{model_name}</a>'
@@ -89,11 +91,13 @@ def make_disclaimer_md():
         '''
     return disclaimer_md
 def make_arena_leaderboard_data(results):
     import pandas as pd
     df = pd.DataFrame(results)
     return df
 def build_leaderboard_tab(score_result_file = 'sorted_score_list.json'):
     with open(score_result_file, "r") as json_file:
         data = json.load(json_file)

 """
     return leaderboard_md
 def make_leaderboard_video_md():
     leaderboard_md = f"""
 # 🏆 K-Sort Arena Leaderboard (Text-to-Video Generation)
 """
     return leaderboard_md
 def model_hyperlink(model_name, link):
     return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{model_name}</a>'
         '''
     return disclaimer_md
 def make_arena_leaderboard_data(results):
     import pandas as pd
     df = pd.DataFrame(results)
     return df
 def build_leaderboard_tab(score_result_file = 'sorted_score_list.json'):
     with open(score_result_file, "r") as json_file:
         data = json.load(json_file)

serve/update_skill.py CHANGED Viewed

@@ -9,9 +9,11 @@ trueskill_env = TrueSkill()
 sys.path.append('../')
 from model.models import IMAGE_GENERATION_MODELS
 ssh_skill_client = None
 sftp_skill_client = None
 def create_ssh_skill_client(server, port, user, password):
     global ssh_skill_client, sftp_skill_client
     ssh_skill_client = paramiko.SSHClient()
@@ -23,32 +25,37 @@ def create_ssh_skill_client(server, port, user, password):
     transport.set_keepalive(60)
     sftp_skill_client = ssh_skill_client.open_sftp()
 def is_connected():
     global ssh_skill_client, sftp_skill_client
     if ssh_skill_client is None or sftp_skill_client is None:
         return False
-    # 检查SSH连接是否正常
     if not ssh_skill_client.get_transport().is_active():
         return False
-    # 检查SFTP连接是否正常
     try:
-        sftp_skill_client.listdir('.')  # 尝试列出根目录
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
@@ -66,6 +73,7 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     with sftp_skill_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_skill_client
     if not is_connected():

 sys.path.append('../')
 from model.models import IMAGE_GENERATION_MODELS
 ssh_skill_client = None
 sftp_skill_client = None
 def create_ssh_skill_client(server, port, user, password):
     global ssh_skill_client, sftp_skill_client
     ssh_skill_client = paramiko.SSHClient()
     transport.set_keepalive(60)
     sftp_skill_client = ssh_skill_client.open_sftp()
 def is_connected():
     global ssh_skill_client, sftp_skill_client
     if ssh_skill_client is None or sftp_skill_client is None:
         return False
     if not ssh_skill_client.get_transport().is_active():
         return False
     try:
+        sftp_skill_client.listdir('.')
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
     with sftp_skill_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_skill_client
     if not is_connected():

serve/update_skill_video.py CHANGED Viewed

@@ -9,9 +9,11 @@ trueskill_env = TrueSkill()
 sys.path.append('../')
 from model.models import VIDEO_GENERATION_MODELS
 ssh_skill_client = None
 sftp_skill_client = None
 def create_ssh_skill_client(server, port, user, password):
     global ssh_skill_client, sftp_skill_client
     ssh_skill_client = paramiko.SSHClient()
@@ -23,32 +25,37 @@ def create_ssh_skill_client(server, port, user, password):
     transport.set_keepalive(60)
     sftp_skill_client = ssh_skill_client.open_sftp()
 def is_connected():
     global ssh_skill_client, sftp_skill_client
     if ssh_skill_client is None or sftp_skill_client is None:
         return False
-    # 检查SSH连接是否正常
     if not ssh_skill_client.get_transport().is_active():
         return False
-    # 检查SFTP连接是否正常
     try:
-        sftp_skill_client.listdir('.')  # 尝试列出根目录
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
@@ -66,6 +73,7 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     with sftp_skill_client.open(SSH_VIDEO_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_skill_client
     if not is_connected():

 sys.path.append('../')
 from model.models import VIDEO_GENERATION_MODELS
 ssh_skill_client = None
 sftp_skill_client = None
 def create_ssh_skill_client(server, port, user, password):
     global ssh_skill_client, sftp_skill_client
     ssh_skill_client = paramiko.SSHClient()
     transport.set_keepalive(60)
     sftp_skill_client = ssh_skill_client.open_sftp()
 def is_connected():
     global ssh_skill_client, sftp_skill_client
     if ssh_skill_client is None or sftp_skill_client is None:
         return False
     if not ssh_skill_client.get_transport().is_active():
         return False
     try:
+        sftp_skill_client.listdir('.')
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
     return ucb
 def update_trueskill(ratings, ranks):
     new_ratings = trueskill_env.rate(ratings, ranks)
     return new_ratings
 def serialize_rating(rating):
     return {'mu': rating.mu, 'sigma': rating.sigma}
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
     with sftp_skill_client.open(SSH_VIDEO_SKILL, 'w') as f:
         f.write(json_data)
 def load_json_via_sftp():
     global sftp_skill_client
     if not is_connected():

serve/upload.py CHANGED Viewed

@@ -9,15 +9,18 @@ import random
 import concurrent.futures
 from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_VIDEO_LOG, SSH_MSCOCO
 ssh_client = None
 sftp_client = None
 sftp_client_imgs = None
 def open_sftp(i=0):
     global ssh_client
     sftp_client = ssh_client.open_sftp()
     return sftp_client
 def create_ssh_client(server, port, user, password):
     global ssh_client, sftp_client, sftp_client_imgs
     ssh_client = paramiko.SSHClient()
@@ -40,22 +43,22 @@ def is_connected():
     global ssh_client, sftp_client
     if ssh_client is None or sftp_client is None:
         return False
-    # 检查SSH连接是否正常
     if not ssh_client.get_transport().is_active():
         return False
-    # 检查SFTP连接是否正常
     try:
-        sftp_client.listdir('.')  # 尝试列出根目录
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def get_image_from_url(image_url):
     response = requests.get(image_url)
     response.raise_for_status()  # success
     return Image.open(io.BytesIO(response.content))
 # def get_random_mscoco_prompt():
 #     global sftp_client
 #     if not is_connected():
@@ -70,6 +73,7 @@ def get_image_from_url(image_url):
 #         print("\n")
 #     return content
 def get_random_mscoco_prompt():
     file_path = './coco_prompt.txt'
@@ -79,6 +83,7 @@ def get_random_mscoco_prompt():
     random_line = random.choice(lines).strip()
     return random_line
 def get_random_video_prompt(root_dir):
     subdirs = [os.path.join(root_dir, d) for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
     if not subdirs:
@@ -96,6 +101,7 @@ def get_random_video_prompt(root_dir):
         raise NotImplementedError
     return selected_dir, prompt
 def get_ssh_random_video_prompt(root_dir, local_dir, model_names):
     def is_directory(sftp, path):
         try:
@@ -150,6 +156,7 @@ def get_ssh_random_video_prompt(root_dir, local_dir, model_names):
     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:
@@ -204,6 +211,7 @@ def get_ssh_random_image_prompt(root_dir, local_dir, model_names):
     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():
@@ -220,6 +228,7 @@ def create_remote_directory(remote_directory, video=False):
         print(f"Directory {remote_directory} created successfully.")
     return log_dir
 def upload_images(i, image_list, output_file_list, sftp_client):
     with sftp_client as sftp:
         if isinstance(image_list[i], str):
@@ -233,7 +242,6 @@ def upload_images(i, image_list, output_file_list, sftp_client):
         print(f"Successfully uploaded image to {output_file_list[i]}")
 def upload_ssh_all(states, output_dir, data, data_path):
     global sftp_client
     global sftp_client_imgs
@@ -246,7 +254,6 @@ def upload_ssh_all(states, output_dir, data, data_path):
         output_file_list.append(output_file)
         image_list.append(states[i].output)
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = [executor.submit(upload_images, i, image_list, output_file_list, sftp_client_imgs[i]) for i in range(len(output_file_list))]
@@ -257,6 +264,7 @@ def upload_ssh_all(states, output_dir, data, data_path):
         print(f"Successfully uploaded JSON data to {data_path}")
     # create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
 def upload_ssh_data(data, data_path):
     global sftp_client
     global sftp_client_imgs

 import concurrent.futures
 from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_VIDEO_LOG, SSH_MSCOCO
 ssh_client = None
 sftp_client = None
 sftp_client_imgs = None
 def open_sftp(i=0):
     global ssh_client
     sftp_client = ssh_client.open_sftp()
     return sftp_client
 def create_ssh_client(server, port, user, password):
     global ssh_client, sftp_client, sftp_client_imgs
     ssh_client = paramiko.SSHClient()
     global ssh_client, sftp_client
     if ssh_client is None or sftp_client is None:
         return False
     if not ssh_client.get_transport().is_active():
         return False
     try:
+        sftp_client.listdir('.')
     except Exception as e:
         print(f"Error checking SFTP connection: {e}")
         return False
     return True
 def get_image_from_url(image_url):
     response = requests.get(image_url)
     response.raise_for_status()  # success
     return Image.open(io.BytesIO(response.content))
 # def get_random_mscoco_prompt():
 #     global sftp_client
 #     if not is_connected():
 #         print("\n")
 #     return content
 def get_random_mscoco_prompt():
     file_path = './coco_prompt.txt'
     random_line = random.choice(lines).strip()
     return random_line
 def get_random_video_prompt(root_dir):
     subdirs = [os.path.join(root_dir, d) for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
     if not subdirs:
         raise NotImplementedError
     return selected_dir, prompt
 def get_ssh_random_video_prompt(root_dir, local_dir, model_names):
     def is_directory(sftp, path):
         try:
     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:
     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():
         print(f"Directory {remote_directory} created successfully.")
     return log_dir
 def upload_images(i, image_list, output_file_list, sftp_client):
     with sftp_client as sftp:
         if isinstance(image_list[i], str):
         print(f"Successfully uploaded image to {output_file_list[i]}")
 def upload_ssh_all(states, output_dir, data, data_path):
     global sftp_client
     global sftp_client_imgs
         output_file_list.append(output_file)
         image_list.append(states[i].output)
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = [executor.submit(upload_images, i, image_list, output_file_list, sftp_client_imgs[i]) for i in range(len(output_file_list))]
         print(f"Successfully uploaded JSON data to {data_path}")
     # create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
 def upload_ssh_data(data, data_path):
     global sftp_client
     global sftp_client_imgs