Upload 7 files

README.md

@@ -4,7 +4,7 @@ emoji: 🤗
 colorFrom: yellow
 colorTo: blue
 sdk: gradio
-sdk_version: 5.4.0
+sdk_version: 5.5.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -13,10 +13,11 @@ css = """
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
     with gr.Column():
         gr.Markdown("# CivitAI to HF🤗 Downloader & Uploader", elem_classes="title")
+        state = gr.State(value={})
         with gr.Accordion("Search Civitai", open=True):
             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_basemodel = gr.CheckboxGroup(label="Base Model", choices=CIVITAI_BASEMODEL, value=[])
             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")
@@ -73,17 +74,18 @@ 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],
-        outputs=[search_civitai_result, search_civitai_desc, search_civitai_submit, search_civitai_query, search_civitai_gallery, search_civitai_result_checkbox, search_civitai_result_info],
+                search_civitai_period, search_civitai_tag, search_civitai_user, search_civitai_limit, search_civitai_page, 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=True,
         show_api=False,
     )
-    search_civitai_result.change(select_civitai_item, [search_civitai_result], [search_civitai_desc, search_civitai_json], queue=False, show_api=False)\
+    search_civitai_result.change(select_civitai_item, [search_civitai_result, state], [search_civitai_desc, search_civitai_json], queue=False, show_api=False)\
     .success(update_civitai_checkbox, [search_civitai_result], [search_civitai_result_checkbox], queue=True, show_api=False)
     search_civitai_result_checkbox.select(from_civitai_checkbox, [search_civitai_result_checkbox], [search_civitai_result], queue=False, show_api=False)
     search_civitai_add.click(add_civitai_item, [search_civitai_result, dl_url], [dl_url], queue=False, show_api=False)
-    search_civitai_select_all.click(select_civitai_all_item, [search_civitai_select_all], [search_civitai_select_all, search_civitai_result], queue=False, show_api=False)
-    search_civitai_gallery.select(update_civitai_selection, [search_civitai_result], [search_civitai_result], queue=False, show_api=False)
+    search_civitai_select_all.click(select_civitai_all_item, [search_civitai_select_all, state], [search_civitai_select_all, search_civitai_result], queue=False, show_api=False)
+    search_civitai_gallery.select(update_civitai_selection, [search_civitai_result, state], [search_civitai_result], queue=False, show_api=False)
 
 demo.queue()
 demo.launch()

civitai_to_hf.py

@@ -6,7 +6,8 @@ import gc
 import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
-from utils import get_token, set_token, is_repo_name, is_repo_exists, get_user_agent, get_download_file, list_uniq, duplicate_hf_repo, HF_SUBFOLDER_NAME
+from utils import (get_token, set_token, is_repo_name, is_repo_exists, get_user_agent, get_download_file,
+                   list_uniq, duplicate_hf_repo, HF_SUBFOLDER_NAME, get_state, set_state)
 import re
 from PIL import Image
 import json
@@ -159,7 +160,7 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
     retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
     session.mount("https://", HTTPAdapter(max_retries=retries))
     try:
-        r = session.get(base_url, params=params, headers=headers, stream=True, timeout=(3.0, 30))
+        r = session.get(base_url, params=params, headers=headers, stream=True, timeout=(7.0, 30))
     except Exception as e:
         print(e)
         return None
@@ -190,40 +191,32 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
         return items
 
 
-civitai_last_results = {}
-civitai_last_choices = [("", "")]
-civitai_last_gallery = []
-
-
-def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", limit=100, page=1, gallery=[]):
-    global civitai_last_results, civitai_last_choices, civitai_last_gallery
-    civitai_last_choices = [("", "")]
-    civitai_last_gallery = []
+def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", limit=100, page=1, gallery=[], state={}):
     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))
     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
+          gr.update(value="", visible=False), gr.update(), gr.update(), gr.update(), gr.update(), results_info, state
     choices = []
     gallery = []
     for item in items:
         base_model_name = "Pony🐴" if item['base_model'] == "Pony" else item['base_model']
-        #name = f"{item['name']} (for {base_model_name} / By: {item['creator']} / Tags: {', '.join(item['tags'])})"
         name = f"{item['name']} (for {base_model_name} / By: {item['creator']})"
         value = item['dl_url']
         choices.append((name, value))
         gallery.append((item['img_url'], name))
         civitai_last_results[value] = item
-    if not choices: return gr.update(choices=[("", "")], value=[], visible=True),\
-          gr.update(value="", visible=False), gr.update(), gr.update(), gr.update(), gr.update(), results_info
-    #result = civitai_last_results.get(choices[0][1], "None") # for debug
-    #md = result['md'] if result else "" # for debug
+    if len(choices) >= 1: results_info = f"{int(len(choices))} items found."
+    else: choices = [("", "")]
     md = ""
-    civitai_last_choices = choices
-    civitai_last_gallery = gallery
-    results_info = f"{int(len(choices))} items found."
+    set_state(state, "civitai_last_choices", choices)
+    set_state(state, "civitai_last_gallery", gallery)
+    set_state(state, "civitai_last_results", civitai_last_results)
     return gr.update(choices=choices, value=[], visible=True), gr.update(value=md, visible=True),\
-          gr.update(), gr.update(), gr.update(value=gallery), gr.update(choices=choices, value=[]), results_info
+          gr.update(), gr.update(), gr.update(value=gallery), gr.update(choices=choices, value=[]), results_info, state
 
 
 def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str=""):
@@ -239,19 +232,19 @@ def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str=""):
     session.mount("https://", HTTPAdapter(max_retries=retries))
     url = base_url + str(dl_url.split("/")[-1])
     try:
-        r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
+        r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
         if not r.ok: return default
         json = dict(r.json()).copy()
         html = ""
         image = ""
         if "modelId" in json.keys():
             url = f"https://civitai.com/models/{json['modelId']}"
-            r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
+            r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
             if not r.ok: return json, html, image
             html = r.text
         if 'images' in json.keys() and len(json["images"]) != 0:
             url = json["images"][0]["url"]
-            r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
+            r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
             if not r.ok: return json, html, image
             image_temp = str(Path(TEMP_DIR, "image" + Path(url.split("/")[-1]).suffix))
             image = str(Path(TEMP_DIR, Path(image_baseurl.split("/")[-1]).stem + ".png"))
@@ -275,7 +268,7 @@ def get_civitai_tag():
     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))
+        r = session.get(url, params=params, headers=headers, stream=True, timeout=(7.0, 15))
         if not r.ok: return default
         j = dict(r.json()).copy()
         if "items" not in j.keys(): return default
@@ -292,11 +285,11 @@ def get_civitai_tag():
         return default
 
 
-def select_civitai_item(results: list[str]):
+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)
-    result = civitai_last_results.get(results[-1], None)
-    md = result['md'] if result else ""
+    result = get_state(state, "civitai_last_results")
+    md = result.get('md', "") if result and isinstance(result, dict) else ""
     #json, html, image = get_civitai_json(results[-1], True) # for debugging
     return gr.update(value=md, visible=True), gr.update(value=json, visible=False)
 
@@ -311,15 +304,17 @@ def add_civitai_item(results: list[str], dl_url: str):
     return gr.update(value=new_url)
 
 
-def select_civitai_all_item(button_name: str):
+def select_civitai_all_item(button_name: str, state: dict):
     if button_name not in ["Select All", "Deselect All"]: return gr.update(value=button_name), gr.Update(visible=True)
+    civitai_last_choices = get_state(state, "civitai_last_choices")
     selected = [t[1] for t in civitai_last_choices if t[1] != ""] if button_name == "Select All" else []
     new_button_name = "Select All" if button_name == "Deselect All" else "Deselect All"
     return gr.update(value=new_button_name), gr.update(value=selected, choices=civitai_last_choices)
 
 
-def update_civitai_selection(evt: gr.SelectData, value: list[str]):
+def update_civitai_selection(evt: gr.SelectData, value: list[str], state: dict):
     try:
+        civitai_last_choices = get_state(state, "civitai_last_choices")
         selected_index = evt.index
         selected = list_uniq([v for v in value if v != ""] + [civitai_last_choices[selected_index][1]])
         return gr.update(value=selected)

utils.py

@@ -7,6 +7,8 @@ import gc
 import re
 import urllib.parse
 import subprocess
+import time
+from typing import Any
 
 
 def get_token():
@@ -24,6 +26,17 @@ def set_token(token):
         print(f"Error: Failed to save token.")
 
 
+def get_state(state: dict, key: str):
+    if key in state.keys(): return state[key]
+    else:
+        print(f"State '{key}' not found.")
+        return None
+
+
+def set_state(state: dict, key: str, value: Any):
+    state[key] = value
+
+
 def get_user_agent():
     return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'
 
@@ -262,7 +275,13 @@ def git_clone(url: str, path: str, pip: bool=False, addcmd: str=""): # requires
         if i != 0: print(f'Error occured at running {cmd}')
 
 
-def run(cmd: str):
+def run(cmd: str, timeout: float=0):
     print(f'Running {cmd} at {Path.cwd()}')
-    i = subprocess.run(cmd, shell=True).returncode
-    if i != 0: print(f'Error occured at running {cmd}')
+    if timeout == 0:
+        i = subprocess.run(cmd, shell=True).returncode
+        if i != 0: print(f'Error occured at running {cmd}')
+    else:
+        p = subprocess.Popen(cmd, shell=True)
+        time.sleep(timeout)
+        p.terminate()
+        print(f'Terminated in {timeout} seconds')