add model list

app.py

@@ -4,16 +4,14 @@ from time import time_ns
 import spaces
 import gradio as gr
 import torch
-from huggingface_hub import Repository
-from llama_cpp import Llama, LLAMA_SPLIT_MODE_NONE
 from transformers import LlamaForCausalLM, LlamaTokenizer
 
 from kgen.generate import tag_gen
 from kgen.metainfo import SPECIAL, TARGET
 
 
-MODEL_PATH = "KBlueLeaf/DanTagGen"
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+MODEL_PATHS = ["KBlueLeaf/DanTagGen-alpha", "KBlueLeaf/DanTagGen-beta"]
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {DEVICE}")
 
 
@@ -113,12 +111,17 @@ masterpiece, newest, absurdres, {rating}"""
 
 
 if __name__ == "__main__":
-    tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
-    text_model = LlamaForCausalLM.from_pretrained(MODEL_PATH)
-    text_model = text_model.eval().half().to(DEVICE)
+    models = {
+        model_path: [
+            LlamaForCausalLM.from_pretrained(model_path).eval().half().to(DEVICE),
+            LlamaTokenizer.from_pretrained(model_path),
+        ]
+        for model_path in MODEL_PATHS
+    }
 
     @spaces.GPU
     def wrapper(
+        model: str,
         rating: str,
         artist: str,
         characters: str,
@@ -132,6 +135,7 @@ if __name__ == "__main__":
         escape_bracket: bool,
         temperature: float = 1.35,
     ):
+        text_model, tokenizer = models[model]
         yield from get_result(
             text_model,
             tokenizer,
@@ -151,7 +155,8 @@ if __name__ == "__main__":
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("""# DanTagGen beta DEMO""")
         with gr.Accordion("Introduction and Instructions"):
-            gr.Markdown("""
+            gr.Markdown(
+                """
 #### What is this:
 DanTagGen(Danbooru Tag Generator) is a LLM model designed for generating Danboou Tags with provided informations.<br>
 It aims to provide user a more convinient way to make prompts for Text2Image model which is trained on Danbooru datasets.
@@ -173,7 +178,8 @@ It aims to provide user a more convinient way to make prompts for Text2Image mod
 #### Notice
 The formated result use same format as what Kohaku-XL Delta used. <br>
 The performance of using the output from this demo for other model is not guaranteed.
-""")
+"""
+            )
         with gr.Row():
             with gr.Column(scale=4):
                 with gr.Row():
@@ -198,9 +204,9 @@ The performance of using the output from this demo for other model is not guaran
                             label="Target length",
                         )
                     with gr.Column(scale=2):
-                        general = gr.TextArea(label="Input your general tags")
+                        general = gr.TextArea(label="Input your general tags", lines=6)
                         black_list = gr.TextArea(
-                            label="tag Black list (seperated by comma)"
+                            label="tag Black list (seperated by comma)", lines=5
                         )
                         with gr.Row():
                             width = gr.Slider(
@@ -229,6 +235,11 @@ The performance of using the output from this demo for other model is not guaran
                                 value=False,
                                 label="Escape bracket",
                             )
+                        model = gr.Dropdown(
+                            list(models.keys()),
+                            value=list(models.keys())[-1],
+                            label="Model",
+                        )
                 submit = gr.Button("Submit")
             with gr.Column(scale=3):
                 formated_result = gr.TextArea(
@@ -239,6 +250,7 @@ The performance of using the output from this demo for other model is not guaran
         submit.click(
             wrapper,
             inputs=[
+                model,
                 rating,
                 artist,
                 characters,