Upload 2 files (#2)

- Upload 2 files (96498f17d1846904bb61dabd4d8a467a77f6ccf2)


Co-authored-by: John Smith <[email protected]>

StableGR.py

@@ -1,357 +1,356 @@
-import gradio as gr
-from huggingface_hub import HfApi, hf_hub_url
-import os
-from pathlib import Path
-import gc
-import requests
-from requests.adapters import HTTPAdapter
-from urllib3.util import Retry
-from utils import (get_user_agent, get_download_file,
-                   list_uniq, list_sub, get_state, set_state)
-import re
-from PIL import Image
-import json
-import pandas as pd
-import tempfile
-import hashlib
-
-
-TEMP_DIR = os.getcwd()
-
-def parse_urls(s):
-    url_pattern = "https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+"
-    try:
-        urls = re.findall(url_pattern, s)
-        return list(urls)
-    except Exception:
-        return []
-
-
-def parse_repos(s):
-    repo_pattern = r'[^\w_\-\.]?([\w_\-\.]+/[\w_\-\.]+)[^\w_\-\.]?'
-    try:
-        s = re.sub("https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+", "", s)
-        repos = re.findall(repo_pattern, s)
-        return list(repos)
-    except Exception:
-        return []
-
-
-def to_urls(l: list[str]):
-    return "\n".join(l)
-
-
-def uniq_urls(s):
-    return to_urls(list_uniq(parse_urls(s) + parse_repos(s)))
-
-
-def generate_image(model_id, prompt, lora_A, num_steps, guidance_scale, sampler, img_width, img_height):
-    model = Model_Diffusers(
-        base_model_id=model_id,
-        task_name='txt2img',
-    )
-
-    image, info_image = model(
-        prompt=prompt,
-        lora_A=lora_A,
-        num_steps=num_steps,
-        guidance_scale=guidance_scale,
-        sampler=sampler,
-        img_width=img_width,
-        img_height=img_height,
-    )
-    return image[0]
-
-
-
-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}")
-    output_filename = get_download_file(download_dir, dl_url, civitai_key)
-    return output_filename
-
-
-def save_civitai_info(dl_url, filename, civitai_key="", progress=gr.Progress(track_tqdm=True)):
-    json_str, html_str, image_path = get_civitai_json(dl_url, True, filename, civitai_key)
-    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 download_civitai(dl_url, civitai_key, hf_token, urls, 
-                     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
-    if not newrepo_id: newrepo_id = os.environ.get("HF_REPO") # default repo to upload
-    if not get_token() or not civitai_key: raise gr.Error("HF write token and Civitai API key is required.")
-    if not urls: urls = []
-    dl_urls = parse_urls(dl_url)
-    remain_urls = dl_urls.copy()
-    try:
-        md = f'### Your repo: [{newrepo_id}]({"https://huggingface.co/datasets/" if repo_type == "dataset" else "https://huggingface.co/"}{newrepo_id})\n'
-        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, civitai_key)
-                urls.append(url)
-                remain_urls.remove(u)
-                md += f"- Uploaded [{str(u)}]({str(u)})\n"
-        dp_repos = parse_repos(dl_url)
-        for r in dp_repos:
-            url = duplicate_hf_repo(r, newrepo_id, "model", repo_type, is_private, HF_SUBFOLDER_NAME[1])
-            if url: urls.append(url)
-        return gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=False)
-    except Exception as e:
-        gr.Info(f"Error occured: {e}")
-        return gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=True)
-    finally:
-        gc.collect()
-
-
-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_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "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,
-                      filetype: list[str] = [], api_key: str = "", progress=gr.Progress(track_tqdm=True)):
-    user_agent = get_user_agent()
-    headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
-    if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}'
-    base_url = 'https://civitai.com/api/v1/models'
-    params = {'sort': sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'}
-    if len(types) != 0: params["types"] = types
-    if query: params["query"] = query
-    if tag: params["tag"] = tag
-    if user: params["username"] = user
-    if page != 0: params["page"] = int(page)
-    session = requests.Session()
-    retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
-    session.mount("https://", HTTPAdapter(max_retries=retries))
-    rs = []
-    try:
-        if page == 0:
-            progress(0, desc="Searching page 1...")
-            print("Searching page 1...")
-            r = session.get(base_url, params=params | {'page': 1}, headers=headers, stream=True, timeout=(7.0, 30))
-            rs.append(r)
-            if r.ok:
-                json = r.json()
-                next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None
-                i = 2
-                while(next_url is not None):
-                    progress(0, desc=f"Searching page {i}...")
-                    print(f"Searching page {i}...")
-                    r = session.get(next_url, headers=headers, stream=True, timeout=(7.0, 30))
-                    rs.append(r)
-                    if r.ok:
-                        json = r.json()
-                        next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None
-                    else: next_url = None
-                    i += 1
-        else:
-            progress(0, desc="Searching page 1...")
-            print("Searching page 1...")
-            r = session.get(base_url, params=params, headers=headers, stream=True, timeout=(7.0, 30))
-            rs.append(r)
-    except requests.exceptions.ConnectTimeout:
-        print("Request timed out.")
-    except Exception as e:
-        print(e)
-    items = []
-    for r in rs:
-        if not r.ok: continue
-        json = r.json()
-        if 'items' not in json: continue
-        for j in json['items']:
-            for model in j['modelVersions']:
-                item = {}
-                if len(allow_model) != 0 and model['baseModel'] not in set(allow_model): continue
-                item['name'] = j['name']
-                item['creator'] = j['creator']['username'] if 'creator' in j.keys() and 'username' in j['creator'].keys() else ""
-                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['md'] = ""
-                if 'images' in model.keys() and len(model["images"]) != 0:
-                    item['img_url'] = model["images"][0]["url"]
-                    item['md'] += f'<img src="{model["images"][0]["url"]}#float" alt="thumbnail" width="150" height="240"><br>'
-                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"]}'''
-                if 'files' in model.keys():
-                    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']
-                    items.append(item)
-    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,
-                   filetype=[], api_key="", 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), filetype, api_key)
-    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 = []
-    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']})"
-        value = item['dl_url']
-        choices.append((name, value))
-        gallery.append((item['img_url'], name))
-        civitai_last_results[value] = item
-    if len(choices) >= 1: results_info = f"{int(len(choices))} items found."
-    else: choices = [("", "")]
-    md = ""
-    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, state
-
-
-def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", api_key=""):
-    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'}
-    if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}'
-    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))
-    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
-        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=(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=(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"))
-            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 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=(7.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], 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), 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 [])
-    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):
-    if "http" not in "".join(results): return gr.update(value=dl_url)
-    new_url = dl_url if dl_url else ""
-    for result in results:
-        if "http" not in result: continue
-        new_url += f"\n{result}" if new_url else f"{result}"
-    new_url = uniq_urls(new_url)
-    return gr.update(value=new_url)
-
-
-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], 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)
-    except Exception:
-        return gr.update()
-
-
-def update_civitai_checkbox(selected: list[str]):
-    return gr.update(value=selected)
-
-
-def from_civitai_checkbox(selected: list[str]):
-    return gr.update(value=selected)
+import gradio as gr
+from huggingface_hub import HfApi, hf_hub_url
+import os
+from pathlib import Path
+import gc
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
+from utils import (get_user_agent, get_download_file,
+                   list_uniq, list_sub, get_state, set_state)
+import re
+from PIL import Image
+import json
+import pandas as pd
+import tempfile
+import hashlib
+from stablepy import Model_Diffusers
+
+
+TEMP_DIR = os.getcwd()
+
+
+def parse_urls(s):
+    url_pattern = "https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+"
+    try:
+        urls = re.findall(url_pattern, s)
+        return list(urls)
+    except Exception:
+        return []
+
+
+def parse_repos(s):
+    repo_pattern = r'[^\w_\-\.]?([\w_\-\.]+/[\w_\-\.]+)[^\w_\-\.]?'
+    try:
+        s = re.sub("https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+", "", s)
+        repos = re.findall(repo_pattern, s)
+        return list(repos)
+    except Exception:
+        return []
+
+
+def to_urls(l: list[str]):
+    return "\n".join(l)
+
+
+def uniq_urls(s):
+    return to_urls(list_uniq(parse_urls(s) + parse_repos(s)))
+
+
+def generate_image(model_id, prompt, lora_A, num_steps, guidance_scale, sampler, img_width, img_height):
+    model = Model_Diffusers(
+        base_model_id=model_id,
+        task_name='txt2img',
+    )
+
+    image, info_image = model(
+        prompt=prompt,
+        lora_A=lora_A,
+        num_steps=num_steps,
+        guidance_scale=guidance_scale,
+        sampler=sampler,
+        img_width=img_width,
+        img_height=img_height,
+    )
+    return image[0]
+
+
+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}")
+    output_filename = get_download_file(download_dir, dl_url, civitai_key)
+    return output_filename
+
+
+def save_civitai_info(dl_url, filename, civitai_key="", progress=gr.Progress(track_tqdm=True)):
+    json_str, html_str, image_path = get_civitai_json(dl_url, True, filename, civitai_key)
+    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 download_civitai(dl_url, civitai_key, hf_token, urls, 
+                     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
+    if not newrepo_id: newrepo_id = os.environ.get("HF_REPO") # default repo to upload
+    if not get_token() or not civitai_key: raise gr.Error("HF write token and Civitai API key is required.")
+    if not urls: urls = []
+    dl_urls = parse_urls(dl_url)
+    remain_urls = dl_urls.copy()
+    try:
+        md = f'### Your repo: [{newrepo_id}]({"https://huggingface.co/datasets/" if repo_type == "dataset" else "https://huggingface.co/"}{newrepo_id})\n'
+        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, civitai_key)
+                urls.append(url)
+                remain_urls.remove(u)
+                md += f"- Uploaded [{str(u)}]({str(u)})\n"
+        dp_repos = parse_repos(dl_url)
+        for r in dp_repos:
+            url = duplicate_hf_repo(r, newrepo_id, "model", repo_type, is_private, HF_SUBFOLDER_NAME[1])
+            if url: urls.append(url)
+        return gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=False)
+    except Exception as e:
+        gr.Info(f"Error occured: {e}")
+        return gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=True)
+    finally:
+        gc.collect()
+
+
+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_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "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,
+                      filetype: list[str] = [], api_key: str = "", progress=gr.Progress(track_tqdm=True)):
+    user_agent = get_user_agent()
+    headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
+    if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}'
+    base_url = 'https://civitai.com/api/v1/models'
+    params = {'sort': sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'}
+    if len(types) != 0: params["types"] = types
+    if query: params["query"] = query
+    if tag: params["tag"] = tag
+    if user: params["username"] = user
+    if page != 0: params["page"] = int(page)
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
+    session.mount("https://", HTTPAdapter(max_retries=retries))
+    rs = []
+    try:
+        if page == 0:
+            progress(0, desc="Searching page 1...")
+            print("Searching page 1...")
+            r = session.get(base_url, params=params | {'page': 1}, headers=headers, stream=True, timeout=(7.0, 30))
+            rs.append(r)
+            if r.ok:
+                json = r.json()
+                next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None
+                i = 2
+                while(next_url is not None):
+                    progress(0, desc=f"Searching page {i}...")
+                    print(f"Searching page {i}...")
+                    r = session.get(next_url, headers=headers, stream=True, timeout=(7.0, 30))
+                    rs.append(r)
+                    if r.ok:
+                        json = r.json()
+                        next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None
+                    else: next_url = None
+                    i += 1
+        else:
+            progress(0, desc="Searching page 1...")
+            print("Searching page 1...")
+            r = session.get(base_url, params=params, headers=headers, stream=True, timeout=(7.0, 30))
+            rs.append(r)
+    except requests.exceptions.ConnectTimeout:
+        print("Request timed out.")
+    except Exception as e:
+        print(e)
+    items = []
+    for r in rs:
+        if not r.ok: continue
+        json = r.json()
+        if 'items' not in json: continue
+        for j in json['items']:
+            for model in j['modelVersions']:
+                item = {}
+                if len(allow_model) != 0 and model['baseModel'] not in set(allow_model): continue
+                item['name'] = j['name']
+                item['creator'] = j['creator']['username'] if 'creator' in j.keys() and 'username' in j['creator'].keys() else ""
+                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['md'] = ""
+                if 'images' in model.keys() and len(model["images"]) != 0:
+                    item['img_url'] = model["images"][0]["url"]
+                    item['md'] += f'<img src="{model["images"][0]["url"]}#float" alt="thumbnail" width="150" height="240"><br>'
+                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"]}'''
+                if 'files' in model.keys():
+                    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']
+                    items.append(item)
+    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,
+                   filetype=[], api_key="", 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), filetype, api_key)
+    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 = []
+    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']})"
+        value = item['dl_url']
+        choices.append((name, value))
+        gallery.append((item['img_url'], name))
+        civitai_last_results[value] = item
+    if len(choices) >= 1: results_info = f"{int(len(choices))} items found."
+    else: choices = [("", "")]
+    md = ""
+    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, state
+
+
+def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", api_key=""):
+    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'}
+    if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}'
+    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))
+    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
+        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=(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=(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"))
+            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 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=(7.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], 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), 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 [])
+    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):
+    if "http" not in "".join(results): return gr.update(value=dl_url)
+    new_url = dl_url if dl_url else ""
+    for result in results:
+        if "http" not in result: continue
+        new_url += f"\n{result}" if new_url else f"{result}"
+    new_url = uniq_urls(new_url)
+    return gr.update(value=new_url)
+
+
+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], 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)
+    except Exception:
+        return gr.update()
+
+
+def update_civitai_checkbox(selected: list[str]):
+    return gr.update(value=selected)
+
+
+def from_civitai_checkbox(selected: list[str]):
+    return gr.update(value=selected)

app.py

@@ -1,120 +1,116 @@
-import gradio as gr
-from stablepy import Model_Diffusers
-
-from StableGR import (generate_image, 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_FILETYPE, download_file)
-
-
-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: right; margin: 20px; }
-"""
-
-
-
-
-with gr.Blocks(fill_width=True, css=css) as demo:
-    with gr.Column():
-        gr.Markdown("# StableGR", elem_classes="title")
-        state = gr.State(value={})
-        with gr.TabItem("Search Civitai"):
-            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")
-                search_civitai_limit = gr.Number(label="Limit", minimum=1, maximum=100, step=1, value=100)
-                search_civitai_page = gr.Number(label="Page", info="If 0, retrieve all pages", minimum=0, maximum=10, step=1, value=1)
-            with gr.Row(equal_height=True):
-                search_civitai_query = gr.Textbox(label="Query", placeholder="flux", 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.TabItem("Results"):
-                with gr.Row():
-                    search_civitai_desc = gr.Markdown(value="", visible=False, elem_classes="desc")
-                    search_civitai_json = gr.JSON(value={}, visible=False)
-                with gr.Row(equal_height=True):
-                    with gr.Column(scale=9):
-                        with gr.TabItem("Select from Gallery"):
-                            search_civitai_gallery = gr.Gallery([], label="Results", allow_preview=False, columns=5, elem_id="gallery", show_share_button=False, interactive=False)
-                        with gr.TabItem("Select by Checkbox"):
-                            search_civitai_result_checkbox = gr.CheckboxGroup(label="", choices=[], value=[])
-                        search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value=[],
-                                                            allow_custom_value=True, visible=True, multiselect=True)
-                        search_civitai_result_info = gr.Markdown("Search result.", elem_classes="info")
-                    with gr.Column(scale=1):
-                        search_civitai_add = gr.Button("Add to download URLs")
-                        search_civitai_select_all = gr.Button("Select All", variant="secondary", size="sm")
-        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.Row():
-            run_base = gr.Button(value="Download Base Model", variant="primary")
-            run_lora = gr.Button(value="Download Lora", variant="primary")
-            uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
-        #urls_md = gr.Markdown("<br><br>", elem_classes="result")
-        urls_remain = gr.Textbox("Remaining URLs", value="", show_copy_button=True, visible=False)
-    with gr.Column():
-        base_model = gr.File(label="Base Models")
-        lora_A = gr.File(label="Lora")
-        with gr.Row():
-            prompt = gr.Textbox(label="Prompt", value="A highly detailed portrait of an underwater city, with towering spires and domes rising up from the ocean floor")
-        
-        num_steps = gr.Slider(label="Number of Steps", minimum=1, maximum=100, value=30, step=1)
-        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20.0, value=7.5, step=0.1)
-        sampler = gr.Dropdown(label="Sampler", choices=["DPM++ 2M", "Euler", "TCD"], value="DPM++ 2M")
-        img_width = gr.Slider(label="Image Width", minimum=64, maximum=2048, value=512, step=64)
-        img_height = gr.Slider(label="Image Height", minimum=64, maximum=2048, value=1024, step=64)
-        generate_button = gr.Button("Generate Image")
-        output_image = gr.Image(label="output")
-
-    gr.on(
-        triggers=[run_base.click],
-        fn=download_file,
-        inputs=[dl_url, civitai_key],
-        outputs=base_model,
-        queue=True,
-    )
-    gr.on(
-        triggers=[run_lora.click],
-        fn=download_file,
-        inputs=[dl_url, civitai_key],
-        outputs=lora_A,
-        queue=True,
-    )
-    gr.on(
-        triggers=[generate_button.click],
-        fn=generate_image,
-        inputs=[prompt, num_steps, guidance_scale, sampler, img_width, img_height],
-        outputs=output_image,
-        queue=True,
-    )
-    gr.on(
-        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_filetype, civitai_key, 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,
-        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)
-    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(debug=True, share=True)
+import gradio as gr
+from StableGR import (generate_image, 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_FILETYPE, download_file)
+
+
+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: right; margin: 20px; }
+"""
+
+
+with gr.Blocks(fill_width=True, css=css) as demo:
+    with gr.Column():
+        gr.Markdown("# StableGR", elem_classes="title")
+        state = gr.State(value={})
+        with gr.TabItem("Search Civitai"):
+            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")
+                search_civitai_limit = gr.Number(label="Limit", minimum=1, maximum=100, step=1, value=100)
+                search_civitai_page = gr.Number(label="Page", info="If 0, retrieve all pages", minimum=0, maximum=10, step=1, value=1)
+            with gr.Row(equal_height=True):
+                search_civitai_query = gr.Textbox(label="Query", placeholder="flux", 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.TabItem("Results"):
+                with gr.Row():
+                    search_civitai_desc = gr.Markdown(value="", visible=False, elem_classes="desc")
+                    search_civitai_json = gr.JSON(value={}, visible=False)
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=9):
+                        with gr.TabItem("Select from Gallery"):
+                            search_civitai_gallery = gr.Gallery([], label="Results", allow_preview=False, columns=5, elem_id="gallery", show_share_button=False, interactive=False)
+                        with gr.TabItem("Select by Checkbox"):
+                            search_civitai_result_checkbox = gr.CheckboxGroup(label="", choices=[], value=[])
+                        search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value=[],
+                                                            allow_custom_value=True, visible=True, multiselect=True)
+                        search_civitai_result_info = gr.Markdown("Search result.", elem_classes="info")
+                    with gr.Column(scale=1):
+                        search_civitai_add = gr.Button("Add to download URLs")
+                        search_civitai_select_all = gr.Button("Select All", variant="secondary", size="sm")
+        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.Row():
+            run_base = gr.Button(value="Download Base Model", variant="primary")
+            run_lora = gr.Button(value="Download Lora", variant="primary")
+            uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
+        #urls_md = gr.Markdown("<br><br>", elem_classes="result")
+        urls_remain = gr.Textbox("Remaining URLs", value="", show_copy_button=True, visible=False)
+    with gr.Column():
+        base_model = gr.File(label="Base Models")
+        lora_A = gr.File(label="Lora")
+        with gr.Row():
+            prompt = gr.Textbox(label="Prompt", value="A highly detailed portrait of an underwater city, with towering spires and domes rising up from the ocean floor")
+        
+        num_steps = gr.Slider(label="Number of Steps", minimum=1, maximum=100, value=30, step=1)
+        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20.0, value=7.5, step=0.1)
+        sampler = gr.Dropdown(label="Sampler", choices=["DPM++ 2M", "Euler", "TCD"], value="DPM++ 2M")
+        img_width = gr.Slider(label="Image Width", minimum=64, maximum=2048, value=512, step=64)
+        img_height = gr.Slider(label="Image Height", minimum=64, maximum=2048, value=1024, step=64)
+        generate_button = gr.Button("Generate Image")
+        output_image = gr.Image(label="output")
+
+    gr.on(
+        triggers=[run_base.click],
+        fn=download_file,
+        inputs=[dl_url, civitai_key],
+        outputs=base_model,
+        queue=True,
+    )
+    gr.on(
+        triggers=[run_lora.click],
+        fn=download_file,
+        inputs=[dl_url, civitai_key],
+        outputs=lora_A,
+        queue=True,
+    )
+    gr.on(
+        triggers=[generate_button.click],
+        fn=generate_image,
+        inputs=[base_model, prompt, lora_A, num_steps, guidance_scale, sampler, img_width, img_height],
+        outputs=output_image,
+        queue=True,
+    )
+    gr.on(
+        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_filetype, civitai_key, 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,
+        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)
+    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(debug=True, share=True)