Upload 3 files

app.py

@@ -3,6 +3,7 @@ from civitai_to_hf import search_civitai, download_civitai, select_civitai_item,
 
 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; }
 """
 
@@ -21,28 +22,34 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
                 search_civitai_tag = gr.Textbox(label="Tag", lines=1)
                 search_civitai_submit = gr.Button("Search on Civitai")
             with gr.Row():
-                search_civitai_json = gr.JSON(value={}, visible=False)
                 search_civitai_desc = gr.Markdown(value="", visible=False)
+                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")
-        dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://civitai.com/api/download/models/28907\n...", value="", lines=2, max_lines=255)
-        civitai_key = gr.Textbox(label="Your Civitai Key", value="", max_lines=1)
-        with gr.Row():
-            hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
-            gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).")
+        with gr.Group():
+            dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://civitai.com/api/download/models/28907\n...", value="", lines=3, max_lines=255)
+            with gr.Column():
+                civitai_key = gr.Textbox(label="Your Civitai Key", value="", max_lines=1)
+                gr.Markdown("Your Civitai API key is available at [https://civitai.com/user/account](https://civitai.com/user/account).", elem_classes="info")
+        with gr.Group():
+            with gr.Row():
+                with gr.Column():
+                    hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
+                    gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).", elem_classes="info")
+                newrepo_id = gr.Textbox(label="Upload repo ID", placeholder="yourid/yourrepo", value="", max_lines=1)
         with gr.Row():
-            newrepo_id = gr.Textbox(label="Upload repo ID", placeholder="yourid/yourrepo", value="", max_lines=1)
             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)
+        run_button = gr.Button(value="Download and Upload", variant="primary")
         uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
-        run_button = gr.Button(value="Download and Upload")
         urls_md = gr.Markdown("<br><br>", elem_classes="result")
     gr.DuplicateButton(value="Duplicate Space")
 
     gr.on(
         triggers=[run_button.click],
         fn=download_civitai,
-        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private],
+        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private, is_info],
         outputs=[uploaded_urls, urls_md],
     )
     gr.on(
@@ -53,7 +60,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_ca
         queue=True,
         show_api=False,
     )
-    search_civitai_result.change(select_civitai_item, [search_civitai_result], [search_civitai_desc], queue=False, 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_add.click(add_civitai_item, [search_civitai_result, dl_url], [dl_url], queue=False, show_api=False)
 
 demo.queue()

civitai_to_hf.py

@@ -8,6 +8,11 @@ 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
 import re
+from PIL import Image
+import json
+
+
+TEMP_DIR = "."
 
 
 def parse_urls(s):
@@ -41,18 +46,61 @@ def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progres
         print(f"Error: Failed to upload to {repo_id}. {e}")
         gr.Warning(f"Error: Failed to upload to {repo_id}. {e}")
         return None
+    finally:
+        Path(filename).unlink()
     return url
 
 
 def download_file(dl_url, civitai_key, progress=gr.Progress(track_tqdm=True)):
-    download_dir = "."
+    download_dir = TEMP_DIR
     progress(0, desc=f"Start downloading... {dl_url}")
     output_filename = get_download_file(download_dir, dl_url, civitai_key)
     return output_filename
 
 
+def save_civitai_info(dl_url, filename, progress=gr.Progress(track_tqdm=True)):
+    json_str, html_str, image_path = get_civitai_json(dl_url, True, filename)
+    if not json_str: return  "", "", ""
+    json_path = str(Path(TEMP_DIR, Path(filename).stem + ".json"))
+    html_path = str(Path(TEMP_DIR, Path(filename).stem + ".html"))
+    try:
+        with open(json_path, 'w') as f:
+            json.dump(json_str, f, indent=2)
+        with open(html_path, mode='w', encoding="utf-8") as f:
+            f.write(html_str)
+        return json_path, html_path, image_path
+    except Exception as e:
+        print(f"Error: Failed to save info file {json_path}, {html_path} {e}")
+        return "", "", ""
+
+
+def upload_info_to_repo(dl_url, filename, repo_id, repo_type, is_private, progress=gr.Progress(track_tqdm=True)):
+    def upload_file(api, filename, repo_id, repo_type, hf_token):
+        if not Path(filename).exists(): return
+        api.upload_file(path_or_fileobj=filename, path_in_repo=Path(filename).name, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id)
+        Path(filename).unlink()
+
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private)
+        progress(0, desc=f"Downloading info... {filename}")
+        json_path, html_path, image_path = save_civitai_info(dl_url, filename)
+        progress(0, desc=f"Start uploading info... {filename} to {repo_id}")
+        if not json_path: return
+        else: upload_file(api, json_path, repo_id, repo_type, hf_token)
+        if html_path: upload_file(api, html_path, repo_id, repo_type, hf_token)
+        if image_path: upload_file(api, image_path, repo_id, repo_type, hf_token)
+        progress(1, desc="Info uploaded.")
+        return
+    except Exception as e:
+        print(f"Error: Failed to upload info to {repo_id}. {e}")
+        gr.Warning(f"Error: Failed to upload info to {repo_id}. {e}")
+        return
+
+
 def download_civitai(dl_url, civitai_key, hf_token, urls, 
-                     newrepo_id, repo_type="model", is_private=True, progress=gr.Progress(track_tqdm=True)):
+                     newrepo_id, repo_type="model", is_private=True, is_info=False, 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
@@ -64,8 +112,9 @@ def download_civitai(dl_url, civitai_key, hf_token, urls,
         file = download_file(u, civitai_key)
         if not Path(file).exists() or not Path(file).is_file(): continue
         url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private)
-        if url: urls.append(url)
-        Path(file).unlink()
+        if url:
+            urls.append(url)
+            if is_info: upload_info_to_repo(u, file, newrepo_id, repo_type, is_private)
     progress(1, desc="Processing...")
     md = f'### Your repo: [{newrepo_id}]({"https://huggingface.co/datasets/" if repo_type == "dataset" else "https://huggingface.co/"}{newrepo_id})\n'
     for u in urls:
@@ -143,11 +192,51 @@ def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIV
           gr.update(visible=True), gr.update(visible=True)
 
 
+def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str=""):
+    if not image_baseurl: image_baseurl = dl_url
+    default = ("", "", "") if is_html else ""
+    if "https://civitai.com/api/download/models/" not in dl_url: return default
+    user_agent = get_user_agent()
+    headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
+    base_url = 'https://civitai.com/api/v1/model-versions/'
+    params = {}
+    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])
+    try:
+        r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.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))
+            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))
+            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"))
+            with open(image_temp, 'wb') as f:
+                f.write(r.content)
+            Image.open(image_temp).convert('RGBA').save(image)
+        return json, html, image
+    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")
     md = result['md'] if result else ""
-    return gr.update(value=md, visible=True)
+    json = {}
+    #json, html, image = get_civitai_json(results[-1], True) # for debugging
+    return gr.update(value=md, visible=True), gr.update(value=json, visible=False)
 
 
 def add_civitai_item(results: list[str], dl_url: str):