Upload 2 files

app.py

@@ -59,6 +59,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
                 newrepo_type = gr.Radio(label="Upload repo type", choices=["model", "dataset"], value="model")
                 is_private = gr.Checkbox(label="Create private repo", value=True)
                 is_info = gr.Checkbox(label="Upload Civitai information files", value=False)
+                is_rename = gr.Checkbox(label="Auto rename", value=True)
         run_button = gr.Button(value="Download and Upload", variant="primary")
         uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
         urls_md = gr.Markdown("<br><br>", elem_classes="result")
@@ -68,7 +69,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
     gr.on(
         triggers=[run_button.click],
         fn=download_civitai,
-        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private, is_info],
+        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private, is_info, is_rename],
         outputs=[uploaded_urls, urls_md, urls_remain],
         queue=True,
     )
@@ -82,7 +83,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
         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)\
+    search_civitai_result.change(select_civitai_item, [search_civitai_result, state], [search_civitai_desc, search_civitai_json, state], 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)
@@ -90,4 +91,4 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
     search_civitai_gallery.select(update_civitai_selection, [search_civitai_result, state], [search_civitai_result], queue=False, show_api=False)
 
 demo.queue()
-demo.launch()
+demo.launch(ssr_mode=False)

civitai_to_hf.py

@@ -7,7 +7,7 @@ import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 from utils import (get_token, set_token, is_repo_exists, get_user_agent, get_download_file,
-                   list_uniq, duplicate_hf_repo, HF_SUBFOLDER_NAME, get_state, set_state)
+                   list_uniq, list_sub, duplicate_hf_repo, HF_SUBFOLDER_NAME, get_state, set_state)
 import re
 from PIL import Image
 import json
@@ -64,6 +64,24 @@ def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progres
     return url
 
 
+def get_safe_filename(filename, repo_id, repo_type):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    new_filename = filename
+    try:
+        i = 1
+        while api.file_exists(repo_id=repo_id, filename=Path(new_filename).name, repo_type=repo_type, token=hf_token):
+            new_filename = str(Path(Path(filename).parent, f"{Path(filename).stem}_{i}{Path(filename).suffix}"))
+            i += 1
+        if filename != new_filename:
+            print(f"{Path(filename).name} is already exists. renaming to {Path(new_filename).name}.")
+            Path(filename).rename(new_filename)
+    except Exception as e:
+        print(f"Error occured when renaming {filename}. {e}")
+    finally:
+        return new_filename
+
+
 def download_file(dl_url, civitai_key, progress=gr.Progress(track_tqdm=True)):
     download_dir = TEMP_DIR
     progress(0, desc=f"Start downloading... {dl_url}")
@@ -113,7 +131,7 @@ def upload_info_to_repo(dl_url, filename, repo_id, repo_type, is_private, progre
 
 
 def download_civitai(dl_url, civitai_key, hf_token, urls, 
-                     newrepo_id, repo_type="model", is_private=True, is_info=False, progress=gr.Progress(track_tqdm=True)):
+                     newrepo_id, repo_type="model", is_private=True, is_info=False, is_rename=True, progress=gr.Progress(track_tqdm=True)):
     if hf_token: set_token(hf_token)
     else: set_token(os.environ.get("HF_TOKEN")) # default huggingface write token
     if not civitai_key: civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key
@@ -127,6 +145,7 @@ def download_civitai(dl_url, civitai_key, hf_token, urls,
         for u in dl_urls:
             file = download_file(u, civitai_key)
             if not Path(file).exists() or not Path(file).is_file(): continue
+            if is_rename: file = get_safe_filename(file, newrepo_id, repo_type)
             url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private)
             if url:
                 if is_info: upload_info_to_repo(u, file, newrepo_id, repo_type, is_private)
@@ -211,7 +230,6 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
                 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['img_url'] = model["images"][0]["url"]
@@ -219,7 +237,14 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
                 else: item['img_url'] = "/home/user/app/null.png"
                 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)
+                if 'files' in model.keys():
+                    for f in model['files']:
+                        i = item.copy()
+                        i['dl_url'] = f['downloadUrl']
+                        items.append(i)
+                else:
+                    item['dl_url'] = model['downloadUrl']
+                    items.append(item)
     return items if len(items) > 0 else None
 
 
@@ -262,7 +287,9 @@ def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str=""):
     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 + str(dl_url.split("/")[-1])
+    model_id = re.sub('https://civitai.com/api/download/models/(\\d+)(?:.+)?', '\\1', dl_url)
+    url = base_url + model_id
+    #url = base_url + str(dl_url.split("/")[-1])
     try:
         r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
         if not r.ok: return default
@@ -321,8 +348,11 @@ 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 = get_state(state, "civitai_last_results")
-    md = result.get('md', "") if result and isinstance(result, dict) else ""
-    return gr.update(value=md, visible=True), gr.update(value=json, visible=False)
+    last_selects = get_state(state, "civitai_last_selects")
+    selects = list_sub(results, last_selects if last_selects else [])
+    md = result.get(selects[-1]).get('md', "") if result and isinstance(result, dict) and len(selects) > 0 else ""
+    set_state(state, "civitai_last_selects", results)
+    return gr.update(value=md, visible=True), gr.update(value=json, visible=False), state
 
 
 def add_civitai_item(results: list[str], dl_url: str):