Upload 12 files

app.py

@@ -505,15 +505,19 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css, delete_cache
                                             lora_md[i] = gr.Markdown(value="", visible=False)
                                             lora_num[i] = gr.Number(i, visible=False)
                             with gr.Accordion("From URL", open=True, visible=True):
+                                with gr.Row():
+                                    lora_search_civitai_basemodel = gr.CheckboxGroup(label="Search LoRA for", choices=["Flux.1 D", "Flux.1 S"], value=["Flux.1 D", "Flux.1 S"])
+                                    lora_search_civitai_sort = gr.Radio(label="Sort", choices=["Highest Rated", "Most Downloaded", "Newest"], value="Highest Rated")
+                                    lora_search_civitai_period = gr.Radio(label="Period", choices=["AllTime", "Year", "Month", "Week", "Day"], value="AllTime")
                                 with gr.Row():
                                     lora_search_civitai_query = gr.Textbox(label="Query", placeholder="flux", lines=1)
+                                    lora_search_civitai_tag = gr.Textbox(label="Tag", lines=1)
                                     lora_search_civitai_submit = gr.Button("Search on Civitai")
-                                    lora_search_civitai_basemodel = gr.CheckboxGroup(label="Search LoRA for", choices=["Flux.1 D", "Flux.1 S"], value=["Flux.1 D", "Flux.1 S"])
                                 with gr.Row():
                                     lora_search_civitai_json = gr.JSON(value={}, visible=False)
                                     lora_search_civitai_desc = gr.Markdown(value="", visible=False)
                                 lora_search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value="", allow_custom_value=True, visible=False)
-                                lora_download_url = gr.Textbox(label="URL", placeholder="http://...my_lora_url.safetensors", lines=1)
+                                lora_download_url = gr.Textbox(label="LoRA URL", placeholder="https://civitai.com/api/download/models/28907", lines=1)
                                 with gr.Row():
                                     lora_download = [None] * num_loras
                                     for i in range(num_loras):
@@ -591,9 +595,9 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css, delete_cache
     prompt_enhance.click(enhance_prompt, [prompt], [prompt], queue=False, show_api=False)
 
     gr.on(
-        triggers=[lora_search_civitai_submit.click, lora_search_civitai_query.submit],
+        triggers=[lora_search_civitai_submit.click, lora_search_civitai_query.submit, lora_search_civitai_tag.submit],
         fn=search_civitai_lora,
-        inputs=[lora_search_civitai_query, lora_search_civitai_basemodel],
+        inputs=[lora_search_civitai_query, lora_search_civitai_basemodel, lora_search_civitai_sort, lora_search_civitai_period, lora_search_civitai_tag],
         outputs=[lora_search_civitai_result, lora_search_civitai_desc, lora_search_civitai_submit, lora_search_civitai_query],
         scroll_to_output=True,
         queue=True,

env.py

@@ -2,7 +2,7 @@ import os
 
 
 CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")
-hf_token = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN")
 hf_read_token = os.environ.get('HF_READ_TOKEN') # only use for private repo
 
 
@@ -67,6 +67,7 @@ HF_MODEL_USER_LIKES = [] # sorted by number of likes
 HF_MODEL_USER_EX = [] # sorted by a special rule
 
 
+
 # - **Download Models**
 download_model_list = [
 ]
@@ -79,6 +80,7 @@ download_vae_list = [
 download_lora_list = [
 ]
 
+DIFFUSERS_FORMAT_LORAS = []
 
 directory_models = 'models'
 os.makedirs(directory_models, exist_ok=True)

flux.py

@@ -11,14 +11,15 @@ warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffuse
 warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
 warnings.filterwarnings(action="ignore", category=FutureWarning, module="transformers")
 from pathlib import Path
-from env import (hf_token, hf_read_token, # to use only for private repos
+from huggingface_hub import HfApi
+from env import (HF_TOKEN, hf_read_token, # to use only for private repos
     CIVITAI_API_KEY, HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2, HF_LORA_ESSENTIAL_PRIVATE_REPO,
     HF_VAE_PRIVATE_REPO, directory_models, directory_loras, directory_vaes,
     download_model_list, download_lora_list, download_vae_list)
 from modutils import (to_list, list_uniq, list_sub, get_lora_model_list, download_private_repo,
     safe_float, escape_lora_basename, to_lora_key, to_lora_path, get_local_model_list,
     get_private_lora_model_lists, get_valid_lora_name, get_valid_lora_path, get_valid_lora_wt,
-    get_lora_info, normalize_prompt_list, get_civitai_info, search_lora_on_civitai)
+    get_lora_info, normalize_prompt_list, get_civitai_info, search_lora_on_civitai, MODEL_TYPE_DICT)
 
 
 def download_things(directory, url, hf_token="", civitai_api_key=""):
@@ -38,7 +39,7 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
         if hf_token:
             os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
         else:
-            os.system (f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+            os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
     elif "civitai.com" in url:
         if "?" in url:
             url = url.split("?")[0]
@@ -94,14 +95,18 @@ vae_model_list = get_model_list(directory_vaes)
 vae_model_list.insert(0, "None")
 
 
+def is_repo_name(s):
+    import re
+    return re.fullmatch(r'^[^/]+?/[^/]+?$', s)
+
+
 def get_t2i_model_info(repo_id: str):
-    from huggingface_hub import HfApi
-    api = HfApi()
+    api = HfApi(token=HF_TOKEN)
     try:
-        if " " in repo_id or not api.repo_exists(repo_id): return ""
-        model = api.model_info(repo_id=repo_id)
+        if not is_repo_name(repo_id): return ""
+        model = api.model_info(repo_id=repo_id, timeout=5.0)
     except Exception as e:
-        print(f"Error: Failed to get {repo_id}'s info. ")
+        print(f"Error: Failed to get {repo_id}'s info.")
         print(e)
         return ""
     if model.private or model.gated: return ""
@@ -109,12 +114,8 @@ def get_t2i_model_info(repo_id: str):
     info = []
     url = f"https://huggingface.co/{repo_id}/"
     if not 'diffusers' in tags: return ""
-    if 'diffusers:FluxPipeline' in tags:
-        info.append("FLUX.1")
-    elif 'diffusers:StableDiffusionXLPipeline' in tags:
-        info.append("SDXL")
-    elif 'diffusers:StableDiffusionPipeline' in tags:
-        info.append("SD1.5")
+    for k, v in MODEL_TYPE_DICT.items():
+        if k in tags: info.append(v)
     if model.card_data and model.card_data.tags:
         info.extend(list_sub(model.card_data.tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']))
     info.append(f"DLs: {model.downloads}")
@@ -246,9 +247,9 @@ def update_loras(prompt, lora, lora_wt):
      gr.update(value=tag, label=label, visible=on), gr.update(value=md, visible=on)
 
 
-def search_civitai_lora(query, base_model):
+def search_civitai_lora(query, base_model, sort="Highest Rated", period="AllTime", tag=""):
     global civitai_lora_last_results
-    items = search_lora_on_civitai(query, base_model)
+    items = search_lora_on_civitai(query, base_model, 100, sort, period, tag)
     if not items: return gr.update(choices=[("", "")], value="", visible=False),\
           gr.update(value="", visible=False), gr.update(visible=True), gr.update(visible=True)
     civitai_lora_last_results = {}

mod.py

@@ -347,7 +347,7 @@ def enhance_prompt(input_prompt):
 
 def save_image(image, savefile, modelname, prompt, height, width, steps, cfg, seed):
     import uuid
-    from PIL import Image, PngImagePlugin
+    from PIL import PngImagePlugin
     import json
     try:
         if savefile is None: savefile = f"{modelname.split('/')[-1]}_{str(uuid.uuid4())}.png"

modutils.py

@@ -4,11 +4,19 @@ import gradio as gr
 from huggingface_hub import HfApi
 import os
 from pathlib import Path
+from PIL import Image
 
 
 from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
-    HF_MODEL_USER_EX, HF_MODEL_USER_LIKES,
-    directory_loras, hf_read_token, hf_token, CIVITAI_API_KEY)
+    HF_MODEL_USER_EX, HF_MODEL_USER_LIKES, DIFFUSERS_FORMAT_LORAS,
+    directory_loras, hf_read_token, HF_TOKEN, CIVITAI_API_KEY)
+
+
+MODEL_TYPE_DICT = {
+    "diffusers:StableDiffusionPipeline": "SD 1.5",
+    "diffusers:StableDiffusionXLPipeline": "SDXL",
+    "diffusers:FluxPipeline": "FLUX",
+}
 
 
 def get_user_agent():
@@ -27,6 +35,11 @@ def list_sub(a, b):
     return [e for e in a if e not in b]
 
 
+def is_repo_name(s):
+    import re
+    return re.fullmatch(r'^[^/]+?/[^/]+?$', s)
+
+
 from translatepy import Translator
 translator = Translator()
 def translate_to_en(input: str):
@@ -64,7 +77,7 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
         if hf_token:
             os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
         else:
-            os.system (f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+            os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
     elif "civitai.com" in url:
         if "?" in url:
             url = url.split("?")[0]
@@ -100,6 +113,23 @@ def safe_float(input):
     return output
 
 
+def save_images(images: list[Image.Image], metadatas: list[str]):
+    from PIL import PngImagePlugin
+    import uuid
+    try:
+        output_images = []
+        for image, metadata in zip(images, metadatas):
+            info = PngImagePlugin.PngInfo()
+            info.add_text("metadata", metadata)
+            savefile = f"{str(uuid.uuid4())}.png"
+            image.save(savefile, "PNG", pnginfo=info)
+            output_images.append(str(Path(savefile).resolve()))
+        return output_images
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e
+
+
 def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
     from datetime import datetime, timezone, timedelta
     progress(0, desc="Updating gallery...")
@@ -209,11 +239,16 @@ def get_model_id_list():
         model_ids.append(model.id) if not model.private else ""
     anime_models = []
     real_models = []
+    anime_models_flux = []
+    real_models_flux = []
     for model in models_ex:
-        if not model.private and not model.gated and "diffusers:FluxPipeline" not in model.tags:
-            anime_models.append(model.id) if "anime" in model.tags else real_models.append(model.id)
+        if not model.private and not model.gated:
+            if "diffusers:FluxPipeline" in model.tags: anime_models_flux.append(model.id) if "anime" in model.tags else real_models_flux.append(model.id)
+            else: anime_models.append(model.id) if "anime" in model.tags else real_models.append(model.id)
     model_ids.extend(anime_models)
     model_ids.extend(real_models)
+    model_ids.extend(anime_models_flux)
+    model_ids.extend(real_models_flux)
     model_id_list = model_ids.copy()
     return model_ids
 
@@ -222,10 +257,10 @@ model_id_list = get_model_id_list()
 
 
 def get_t2i_model_info(repo_id: str):
-    api = HfApi()
+    api = HfApi(token=HF_TOKEN)
     try:
-        if " " in repo_id or not api.repo_exists(repo_id): return ""
-        model = api.model_info(repo_id=repo_id)
+        if not is_repo_name(repo_id): return ""
+        model = api.model_info(repo_id=repo_id, timeout=5.0)
     except Exception as e:
         print(f"Error: Failed to get {repo_id}'s info.")
         print(e)
@@ -235,9 +270,8 @@ def get_t2i_model_info(repo_id: str):
     info = []
     url = f"https://huggingface.co/{repo_id}/"
     if not 'diffusers' in tags: return ""
-    if 'diffusers:FluxPipeline' in tags: info.append("FLUX.1")
-    elif 'diffusers:StableDiffusionXLPipeline' in tags: info.append("SDXL")
-    elif 'diffusers:StableDiffusionPipeline' in tags: info.append("SD1.5")
+    for k, v in MODEL_TYPE_DICT.items():
+        if k in tags: info.append(v)
     if model.card_data and model.card_data.tags:
         info.extend(list_sub(model.card_data.tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']))
     info.append(f"DLs: {model.downloads}")
@@ -262,12 +296,8 @@ def get_tupled_model_list(model_list):
         tags = model.tags
         info = []
         if not 'diffusers' in tags: continue
-        if 'diffusers:FluxPipeline' in tags:
-            info.append("FLUX.1")
-        if 'diffusers:StableDiffusionXLPipeline' in tags:
-            info.append("SDXL")
-        elif 'diffusers:StableDiffusionPipeline' in tags:
-            info.append("SD1.5")
+        for k, v in MODEL_TYPE_DICT.items():
+            if k in tags: info.append(v)
         if model.card_data and model.card_data.tags:
             info.extend(list_sub(model.card_data.tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']))
         if "pony" in info:
@@ -351,7 +381,7 @@ def get_civitai_info(path):
 
 
 def get_lora_model_list():
-    loras = list_uniq(get_private_lora_model_lists() + get_local_model_list(directory_loras))
+    loras = list_uniq(get_private_lora_model_lists() + get_local_model_list(directory_loras) + DIFFUSERS_FORMAT_LORAS)
     loras.insert(0, "None")
     loras.insert(0, "")
     return loras
@@ -408,7 +438,7 @@ def download_lora(dl_urls: str):
     for url in [url.strip() for url in dl_urls.split(',')]:
         local_path = f"{directory_loras}/{url.split('/')[-1]}"
         if not Path(local_path).exists():
-            download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+            download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
             urls.append(url)
     after = get_local_model_list(directory_loras)
     new_files = list_sub(after, before)
@@ -460,7 +490,7 @@ def download_my_lora(dl_urls: str, lora1: str, lora2: str, lora3: str, lora4: st
         gr.update(value=lora4, choices=choices), gr.update(value=lora5, choices=choices)
 
 
-def get_valid_lora_name(query: str):
+def get_valid_lora_name(query: str, model_name: str):
     path = "None"
     if not query or query == "None": return "None"
     if to_lora_key(query) in loras_dict.keys(): return query
@@ -474,7 +504,7 @@ def get_valid_lora_name(query: str):
         dl_file = download_lora(query)
         if dl_file and Path(dl_file).exists(): return dl_file
     else:
-        dl_file = find_similar_lora(query)
+        dl_file = find_similar_lora(query, model_name)
         if dl_file and Path(dl_file).exists(): return dl_file
     return "None"
 
@@ -498,14 +528,14 @@ def get_valid_lora_wt(prompt: str, lora_path: str, lora_wt: float):
     return wt
 
 
-def set_prompt_loras(prompt, prompt_syntax, lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt):
+def set_prompt_loras(prompt, prompt_syntax, model_name, lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt):
     import re
     if not "Classic" in str(prompt_syntax):  return lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt
-    lora1 = get_valid_lora_name(lora1)
-    lora2 = get_valid_lora_name(lora2)
-    lora3 = get_valid_lora_name(lora3)
-    lora4 = get_valid_lora_name(lora4)
-    lora5 = get_valid_lora_name(lora5)
+    lora1 = get_valid_lora_name(lora1, model_name)
+    lora2 = get_valid_lora_name(lora2, model_name)
+    lora3 = get_valid_lora_name(lora3, model_name)
+    lora4 = get_valid_lora_name(lora4, model_name)
+    lora5 = get_valid_lora_name(lora5, model_name)
     if not "<lora" in prompt: return lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt
     lora1_wt = get_valid_lora_wt(prompt, lora1, lora1_wt)
     lora2_wt = get_valid_lora_wt(prompt, lora2, lora2_wt)
@@ -670,7 +700,7 @@ def get_my_lora(link_url):
     before = get_local_model_list(directory_loras)
     for url in [url.strip() for url in link_url.split(',')]:
         if not Path(f"{directory_loras}/{url.split('/')[-1]}").exists():
-            download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+            download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
     after = get_local_model_list(directory_loras)
     new_files = list_sub(after, before)
     for file in new_files:
@@ -727,8 +757,7 @@ def move_file_lora(filepaths):
 
 
 def get_civitai_info(path):
-    global civitai_not_exists_list
-    global loras_url_to_path_dict
+    global civitai_not_exists_list, loras_url_to_path_dict
     import requests
     from requests.adapters import HTTPAdapter
     from urllib3.util import Retry
@@ -768,16 +797,18 @@ def get_civitai_info(path):
         return items
 
 
-def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1.0"], limit: int = 100):
+def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1.0"], limit: int = 100,
+                           sort: str = "Highest Rated", period: str = "AllTime", tag: str = ""):
     import requests
     from requests.adapters import HTTPAdapter
     from urllib3.util import Retry
-    if not query: return None
+    
     user_agent = get_user_agent()
     headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
     base_url = 'https://civitai.com/api/v1/models'
-    params = {'query': query, 'types': ['LORA'], 'sort': 'Highest Rated', 'period': 'AllTime',
-              'nsfw': 'true', 'supportsGeneration ': 'true', 'limit': limit}
+    params = {'types': ['LORA'], 'sort': sort, 'period': period, 'limit': limit, 'nsfw': 'true'}
+    if query: params["query"] = query
+    if tag: params["tag"] = tag
     session = requests.Session()
     retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
     session.mount("https://", HTTPAdapter(max_retries=retries))
@@ -806,9 +837,9 @@ def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1
         return items
 
 
-def search_civitai_lora(query, base_model):
+def search_civitai_lora(query, base_model, sort="Highest Rated", period="AllTime", tag=""):
     global civitai_lora_last_results
-    items = search_lora_on_civitai(query, base_model)
+    items = search_lora_on_civitai(query, base_model, 100, sort, period, tag)
     if not items: return gr.update(choices=[("", "")], value="", visible=False),\
           gr.update(value="", visible=False), gr.update(visible=True), gr.update(visible=True)
     civitai_lora_last_results = {}
@@ -834,7 +865,27 @@ def select_civitai_lora(search_result):
     return gr.update(value=search_result), gr.update(value=md, visible=True)
 
 
-def find_similar_lora(q: str):
+LORA_BASE_MODEL_DICT = {
+    "diffusers:StableDiffusionPipeline": ["SD 1.5"],
+    "diffusers:StableDiffusionXLPipeline": ["Pony", "SDXL 1.0"],
+    "diffusers:FluxPipeline": ["Flux.1 D", "Flux.1 S"],
+}
+
+
+def get_lora_base_model(model_name: str):
+    api = HfApi(token=HF_TOKEN)
+    default = ["Pony", "SDXL 1.0"]
+    try:
+        model = api.model_info(repo_id=model_name, timeout=5.0)
+        tags = model.tags
+        for tag in tags:
+            if tag in LORA_BASE_MODEL_DICT.keys(): return LORA_BASE_MODEL_DICT.get(tag, default)
+    except Exception:
+        return default
+    return default
+
+
+def find_similar_lora(q: str, model_name: str):
     from rapidfuzz.process import extractOne
     from rapidfuzz.utils import default_process
     query = to_lora_key(q)
@@ -857,7 +908,7 @@ def find_similar_lora(q: str):
     print(f"Finding <lora:{query}:...> on Civitai...")
     civitai_query = Path(query).stem if Path(query).is_file() else query
     civitai_query = civitai_query.replace("_", " ").replace("-", " ")
-    base_model = ["Pony", "SDXL 1.0"]
+    base_model = get_lora_base_model(model_name)
     items = search_lora_on_civitai(civitai_query, base_model, 1)
     if items:
         item = items[0]
@@ -1219,12 +1270,12 @@ def set_textual_inversion_prompt(textual_inversion_gui, prompt_gui, neg_prompt_g
 
 def get_model_pipeline(repo_id: str):
     from huggingface_hub import HfApi
-    api = HfApi()
+    api = HfApi(token=HF_TOKEN)
     default = "StableDiffusionPipeline"
     try:
-        if " " in repo_id or not api.repo_exists(repo_id): return default
-        model = api.model_info(repo_id=repo_id)
-    except Exception as e:
+        if not is_repo_name(repo_id): return default
+        model = api.model_info(repo_id=repo_id, timeout=5.0)
+    except Exception:
         return default
     if model.private or model.gated: return default
     tags = model.tags