Upload 2 files

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
-from civitai_to_hf import search_civitai, download_civitai, select_civitai_item, add_civitai_item, CIVITAI_TYPE, CIVITAI_BASEMODEL, CIVITAI_SORT, CIVITAI_PERIOD
+from civitai_to_hf import (search_civitai, download_civitai, select_civitai_item, add_civitai_item, get_civitai_tag,
+                           CIVITAI_TYPE, CIVITAI_BASEMODEL, CIVITAI_SORT, CIVITAI_PERIOD)
 
 css = """
 .title { font-size: 3em; align-items: center; text-align: center; }
 .info { align-items: center; text-align: center; }
 .block.result { margin: 1em 0; padding: 1em; box-shadow: 0 0 3px 3px #664422, 0 0 3px 2px #664422 inset; border-radius: 6px; background: #665544; }
+.desc [src$='#float'] { float: left; margin: 10px; }
 """
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
@@ -19,10 +21,11 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
                 search_civitai_period = gr.Radio(label="Period", choices=CIVITAI_PERIOD, value="Month")
             with gr.Row():
                 search_civitai_query = gr.Textbox(label="Query", placeholder="flux", lines=1)
-                search_civitai_tag = gr.Textbox(label="Tag", lines=1)
+                search_civitai_tag = gr.Dropdown(label="Tag", choices=get_civitai_tag(), value=get_civitai_tag()[0], allow_custom_value=True)
+                search_civitai_user = gr.Textbox(label="Username", lines=1)
                 search_civitai_submit = gr.Button("Search on Civitai")
             with gr.Row():
-                search_civitai_desc = gr.Markdown(value="", visible=False)
+                search_civitai_desc = gr.Markdown(value="", visible=False, elem_classes="desc")
                 search_civitai_json = gr.JSON(value={}, visible=False)
             search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value="", allow_custom_value=True, visible=False, multiselect=True)
             search_civitai_add = gr.Button("Add to download URLs")
@@ -53,9 +56,9 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
         outputs=[uploaded_urls, urls_md],
     )
     gr.on(
-        triggers=[search_civitai_submit.click, search_civitai_query.submit, search_civitai_tag.submit],
+        triggers=[search_civitai_submit.click, search_civitai_query.submit, search_civitai_user.submit],
         fn=search_civitai,
-        inputs=[search_civitai_query, search_civitai_type, search_civitai_basemodel, search_civitai_sort, search_civitai_period, search_civitai_tag],
+        inputs=[search_civitai_query, search_civitai_type, search_civitai_basemodel, search_civitai_sort, search_civitai_period, search_civitai_tag, search_civitai_user],
         outputs=[search_civitai_result, search_civitai_desc, search_civitai_submit, search_civitai_query],
         queue=True,
         show_api=False,

civitai_to_hf.py

@@ -10,6 +10,7 @@ from utils import get_token, set_token, is_repo_exists, get_user_agent, get_down
 import re
 from PIL import Image
 import json
+import pandas as pd
 
 
 TEMP_DIR = "."
@@ -130,7 +131,7 @@ CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"]
 
 
 def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [], limit: int = 100,
-                           sort: str = "Highest Rated", period: str = "AllTime", tag: str = ""):
+                           sort: str = "Highest Rated", period: str = "AllTime", tag: str = "", user: str = ""):
     user_agent = get_user_agent()
     headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
     base_url = 'https://civitai.com/api/v1/models'
@@ -138,6 +139,7 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
     if len(types) != 0: params["types"] = types
     if query: params["query"] = query
     if tag: params["tag"] = tag
+    if user: params["username"] = user
     session = requests.Session()
     retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
     session.mount("https://", HTTPAdapter(max_retries=retries))
@@ -160,10 +162,13 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
                 item['tags'] = j['tags'] if 'tags' in j.keys() else []
                 item['model_name'] = model['name'] if 'name' in model.keys() else ""
                 item['base_model'] = model['baseModel'] if 'baseModel' in model.keys() else ""
+                item['description'] = model['description'] if 'description' in model.keys() else ""
                 item['dl_url'] = model['downloadUrl']
+                item['md'] = ""
                 if 'images' in model.keys() and len(model["images"]) != 0:
-                    item['md'] = f'<img src="{model["images"][0]["url"]}" alt="thumbnail" width="150" height="240"><br>[Model URL](https://civitai.com/models/{j["id"]})'
-                else: item['md'] = f'[Model URL](https://civitai.com/models/{j["id"]})'
+                    item['md'] += f'<img src="{model["images"][0]["url"]}#float" alt="thumbnail" width="150" height="240"><br>'
+                item['md'] += f'''Model URL: [https://civitai.com/models/{j["id"]}](https://civitai.com/models/{j["id"]})<br>Model Name: {item["name"]}<br>
+                    Creator: {item["creator"]}<br>Tags: {", ".join(item["tags"])}<br>Base Model: {item["base_model"]}<br>Description: {item["description"]}'''
                 items.append(item)
         return items
 
@@ -171,9 +176,9 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
 civitai_last_results = {}
 
 
-def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag=""):
+def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user=""):
     global civitai_last_results
-    items = search_on_civitai(query, types, base_model, 100, sort, period, tag)
+    items = search_on_civitai(query, types, base_model, 100, sort, period, tag, user)
     if not items: return gr.update(choices=[("", "")], value="", visible=False),\
           gr.update(value="", visible=False), gr.update(visible=True), gr.update(visible=True)
     civitai_last_results = {}
@@ -230,6 +235,34 @@ def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str=""):
         return default
 
 
+def get_civitai_tag():
+    default = [""]
+    user_agent = get_user_agent()
+    headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
+    base_url = 'https://civitai.com/api/v1/tags'
+    params = {'limit': 200}
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
+    session.mount("https://", HTTPAdapter(max_retries=retries))
+    url = base_url
+    try:
+        r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
+        if not r.ok: return default
+        j = dict(r.json()).copy()
+        if "items" not in j.keys(): return default
+        items = []
+        for item in j["items"]:
+            items.append([str(item.get("name", "")), int(item.get("modelCount", 0))])
+        df = pd.DataFrame(items)
+        df.sort_values(1, ascending=False)
+        tags = df.values.tolist()
+        tags = [""] + [l[0] for l in tags]
+        return tags
+    except Exception as e:
+        print(e)
+        return default
+
+
 def select_civitai_item(results: list[str]):
     if "http" not in "".join(results): return gr.update(value="None", visible=True)
     result = civitai_last_results.get(results[-1], "None")