Update app.py

app.py

@@ -7,6 +7,31 @@ from PIL import Image
 from huggingface_hub import InferenceApi, InferenceClient
 from datasets import load_dataset
 import pandas as pd
+import re
+
+def rank_score(repo_str):
+    p_list = re.findall(r"[-_xlv\d]+" ,repo_str.split("/")[-1])
+    xl_in_str = any(map(lambda x: "xl" in x, p_list))
+    v_in_str = any(map(lambda x: "v" in x and 
+                       any(map(lambda y: 
+                               any(map(lambda z: y.startswith(z), "0123456789"))
+                               ,x.split("v")))
+                       , p_list))
+    stable_in_str = repo_str.split("/")[-1].lower().startswith("stable")
+    score = sum(map(lambda t2: t2[0] * t2[1] ,(zip(*[[stable_in_str, xl_in_str, v_in_str], [1000, 100, 10]]))))
+    #return p_list, xl_in_str, v_in_str, stable_in_str, score
+    return score
+
+def shorten_by(repo_list, by = None):
+    if by == "user":
+        return sorted(
+            pd.DataFrame(pd.Series(repo_list).map(lambda x: (x.split("/")[0], x)).values.tolist()).groupby(0)[1].apply(list).map(lambda x: 
+                                                                                                                                    sorted(x, key = rank_score, reverse = True)[0]).values.tolist(),
+            key = rank_score, reverse = True
+        )
+    if by == "model":
+        return sorted(repo_list, key = lambda x: rank_score(x), reverse = True)
+    return repo_list
 
 '''
 dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts")
@@ -21,26 +46,25 @@ def get_samples():
     prompt_list = prompt_df.sample(n = 10)["Prompt"].map(lambda x: x).values.tolist()
     return prompt_list
 
-def update_models():
+def update_models(models_rank_by = "model"):
     client = InferenceClient()
     models = client.list_deployed_models()
     list_models = models["text-to-image"]
+    if hasattr(models_rank_by, "value"):
+        list_models = shorten_by(list_models, models_rank_by.value)
+    else:
+        list_models = shorten_by(list_models, models_rank_by)
     return gr.Dropdown.update(choices=list_models) 
 
 def update_prompts():
     return gr.Dropdown.update(choices=get_samples())
 
-def get_params(request: gr.Request):
+def get_params(request: gr.Request, models_rank_by):
     params = request.query_params
     ip = request.client.host
     req = {"params": params,
           "ip": ip}
-    return update_models(), update_prompts()
-
-client = InferenceClient()
-models = client.list_deployed_models()
-list_models = models["text-to-image"]
-list_prompts = get_samples()
+    return update_models(models_rank_by), update_prompts()
 
 '''
 list_models = [
@@ -239,6 +263,12 @@ with gr.Blocks(css=css) as demo:
     with gr.Row(elem_id="prompt-container"):
         with gr.Column():
             btn_refresh = gr.Button(value="Click to get current deployed models and newly Prompt candidates")
+            models_rank_by = gr.Radio(choices=["model", "user"],
+                            value="model", label="Models ranked by", elem_id="rank_radio")
+            
+            list_models = update_models(models_rank_by)
+            list_prompts = get_samples()
+            
             #btn_refresh.click(None, js="window.location.reload()")
             current_model = gr.Dropdown(label="Current Model", choices=list_models, value=DEFAULT_MODEL,
                                        info = "default model: {}".format(DEFAULT_MODEL)
@@ -276,9 +306,11 @@ with gr.Blocks(css=css) as demo:
             
     text_button.click(generate_txt2img, inputs=[current_model, text_prompt, negative_prompt, image_style], outputs=image_output)
     select_button.click(generate_txt2img, inputs=[current_model, select_prompt, negative_prompt, image_style], outputs=image_output)
-    btn_refresh.click(update_models, None, current_model)
+    btn_refresh.click(update_models, models_rank_by, current_model)
     btn_refresh.click(update_prompts, None, select_prompt)
+
+    models_rank_by.change(update_models, models_rank_by, current_model)
     
-    demo.load(get_params, None, [current_model, select_prompt])
+    demo.load(get_params, models_rank_by, [current_model, select_prompt])
 
 demo.launch(show_api=False)