GenAI-Arena / model /model_registry.py
tianleliphoebe's picture
Reinitialize
e368cec
raw
history blame
5.49 kB
from collections import namedtuple
from typing import List
ModelInfo = namedtuple("ModelInfo", ["simple_name", "link", "description"])
model_info = {}
def register_model_info(
full_names: List[str], simple_name: str, link: str, description: str
):
info = ModelInfo(simple_name, link, description)
for full_name in full_names:
model_info[full_name] = info
def get_model_info(name: str) -> ModelInfo:
if name in model_info:
return model_info[name]
else:
# To fix this, please use `register_model_info` to register your model
return ModelInfo(
name, "", "Register the description at fastchat/model/model_registry.py"
)
def get_model_description_md(model_list):
model_description_md = """
| | | |
| ---- | ---- | ---- |
"""
ct = 0
visited = set()
for i, name in enumerate(model_list):
minfo = get_model_info(name)
if minfo.simple_name in visited:
continue
visited.add(minfo.simple_name)
one_model_md = f"[{minfo.simple_name}]({minfo.link}): {minfo.description}"
if ct % 3 == 0:
model_description_md += "|"
model_description_md += f" {one_model_md} |"
if ct % 3 == 2:
model_description_md += "\n"
ct += 1
return model_description_md
# regist image generation models
register_model_info(
["imagenhub_LCM_generation"],
"LCM",
"https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7",
"Latent Consistency Models.",
)
register_model_info(
["imagenhub_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 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"],
"SDXLTurbo",
"https://huggingface.co/stabilityai/sdxl-turbo",
"SDXL-Turbo is a fast generative text-to-image model.",
)
register_model_info(
["imagenhub_SDXL_generation"],
"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_SDXLLightning_generation"],
"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"],
"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_SDEdit_edition"],
# "SDEdit",
# "",
# "xxx",
# )
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(
["fal_stable_cascade"],
"StableCascade",
"https://fal.ai/models/stable-cascade/api",
"StableCascade is a generative model that can generate high-quality images from text prompts.",
)
models = ['imagenhub_LCM_generation','imagenhub_SDXLTurbo_generation','imagenhub_SDXL_generation',
'imagenhub_OpenJourney_generation','imagenhub_PixArtAlpha_generation','imagenhub_SDXLLightning_generation',
'imagenhub_StableCascade_generation','imagenhub_PlaygroundV2_generation', 'fal_Playground-v25_generation', 'fal_stable_cascade',
'imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition']