Upload 3 files

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from civitai_to_hf import (search_civitai, download_civitai, select_civitai_item, add_civitai_item, get_civitai_tag, select_civitai_all_item,
                            update_civitai_selection, update_civitai_checkbox, from_civitai_checkbox,
-                           CIVITAI_TYPE, CIVITAI_BASEMODEL, CIVITAI_SORT, CIVITAI_PERIOD)
+                           CIVITAI_TYPE, CIVITAI_BASEMODEL, CIVITAI_SORT, CIVITAI_PERIOD, CIVITAI_FILETYPE)
 
 css = """
 .title { font-size: 3em; align-items: center; text-align: center; }
@@ -18,6 +18,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
             with gr.Row():
                 search_civitai_type = gr.CheckboxGroup(label="Type", choices=CIVITAI_TYPE, value=["Checkpoint", "LORA"])
                 search_civitai_basemodel = gr.CheckboxGroup(label="Base Model", choices=CIVITAI_BASEMODEL, value=[])
+                search_civitai_filetype = gr.CheckboxGroup(label="File type", choices=CIVITAI_FILETYPE, value=["Model"])
             with gr.Row():
                 search_civitai_sort = gr.Radio(label="Sort", choices=CIVITAI_SORT, value=CIVITAI_SORT[0])
                 search_civitai_period = gr.Radio(label="Period", choices=CIVITAI_PERIOD, value="Month")
@@ -77,7 +78,8 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
         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, search_civitai_user, search_civitai_limit, search_civitai_page, search_civitai_gallery, state],
+                search_civitai_period, search_civitai_tag, search_civitai_user, search_civitai_limit,
+                search_civitai_page, search_civitai_filetype, search_civitai_gallery, state],
         outputs=[search_civitai_result, search_civitai_desc, search_civitai_submit, search_civitai_query, search_civitai_gallery,
                  search_civitai_result_checkbox, search_civitai_result_info, state],
         queue=False,

civitai_to_hf.py

@@ -185,13 +185,15 @@ def download_civitai(dl_url, civitai_key, hf_token, urls,
 
 CIVITAI_TYPE = ["Checkpoint", "TextualInversion", "Hypernetwork", "AestheticGradient", "LORA", "LoCon", "DoRA",
                 "Controlnet", "Upscaler", "MotionModule", "VAE", "Poses", "Wildcards", "Workflows", "Other"]
+CIVITAI_FILETYPE = ["Model", "VAE", "Config", "Training Data"]
 CIVITAI_BASEMODEL = ["Pony", "Illustrious", "SDXL 1.0", "SD 1.5", "Flux.1 D", "Flux.1 S", "SD 3.5"]
 CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Newest"]
 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 = "", user: str = "", page: int = 1, progress=gr.Progress(track_tqdm=True)):
+                      sort: str = "Highest Rated", period: str = "AllTime", tag: str = "", user: str = "", page: int = 1,
+                      filetype: list[str] = [], progress=gr.Progress(track_tqdm=True)):
     user_agent = get_user_agent()
     headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
     base_url = 'https://civitai.com/api/v1/models'
@@ -260,6 +262,7 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
                     for f in model['files']:
                         i = item.copy()
                         i['dl_url'] = f['downloadUrl']
+                        if len(filetype) != 0 and f['type'] not in set(filetype): continue
                         items.append(i)
                 else:
                     item['dl_url'] = model['downloadUrl']
@@ -267,13 +270,14 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
     return items if len(items) > 0 else None
 
 
-def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", limit=100, page=1, gallery=[], state={}, progress=gr.Progress(track_tqdm=True)):
+def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", limit=100, page=1,
+                   filetype=[], gallery=[], state={}, progress=gr.Progress(track_tqdm=True)):
     civitai_last_results = {}
     set_state(state, "civitai_last_choices", [("", "")])
     set_state(state, "civitai_last_gallery", [])
     set_state(state, "civitai_last_results", civitai_last_results)
     results_info = "No item found."
-    items = search_on_civitai(query, types, base_model, int(limit), sort, period, tag, user, int(page))
+    items = search_on_civitai(query, types, base_model, int(limit), sort, period, tag, user, int(page), filetype)
     if not items: return gr.update(choices=[("", "")], value=[], visible=True),\
           gr.update(value="", visible=False), gr.update(), gr.update(), gr.update(), gr.update(), results_info, state
     choices = []
@@ -365,7 +369,7 @@ def get_civitai_tag():
 
 def select_civitai_item(results: list[str], state: dict):
     json = {}
-    if "http" not in "".join(results) or len(results) == 0: return gr.update(value="", visible=True), gr.update(value=json, visible=False)
+    if "http" not in "".join(results) or len(results) == 0: return gr.update(value="", visible=True), gr.update(value=json, visible=False), state
     result = get_state(state, "civitai_last_results")
     last_selects = get_state(state, "civitai_last_selects")
     selects = list_sub(results, last_selects if last_selects else [])