Upload 4 files

app.py

@@ -21,7 +21,8 @@ from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists
                  get_control_union_mode, set_control_union_mode, get_control_params, translate_to_en)
 from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json,
                   download_my_lora, get_all_lora_tupled_list, apply_lora_prompt,
-                  update_loras, get_t2i_model_info)
+                  update_loras)
+from modutils import get_t2i_model_info, get_download_file
 from tagger.tagger import predict_tags_wd, compose_prompt_to_copy
 from tagger.fl2flux import predict_tags_fl2_flux
 
@@ -111,6 +112,11 @@ def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, prog
 
 change_base_model.zerogpu = True
 
+def download_file_mod(url, directory=os.getcwd()):
+    path = get_download_file(directory, url)
+    if not path: raise Exception(f"Download error: {url}")
+    return path
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -130,7 +136,7 @@ class calculateDuration:
 def download_file(url, directory=None):
     if directory is None:
         directory = os.getcwd()  # Use current working directory if not specified
-    
+
     # Get the filename from the URL
     filename = url.split('/')[-1]
     
@@ -146,7 +152,7 @@ def download_file(url, directory=None):
         file.write(response.content)
     
     return filepath
-            
+
 def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, height):
     selected_index = evt.index
     selected_indices = selected_indices or []
@@ -243,7 +249,8 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
             existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
             if existing_item_index is None:
                 if repo.endswith(".safetensors") and repo.startswith("http"):
-                    repo = download_file(repo)
+                    #repo = download_file(repo)
+                    repo = download_file_mod(repo)
                 new_item = {
                     "image": image if image else "/home/user/app/custom.png",
                     "title": title,
@@ -504,14 +511,14 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             print(f"Lora Path: {lora_path}")
             if image_input is not None:
                 if weight_name:
-                    pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name)
+                    pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
                 else:
-                    pipe_i2i.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name)
+                    pipe_i2i.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
             else:
                 if weight_name:
-                    pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name)
+                    pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
                 else:
-                    pipe.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name)
+                    pipe.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
         print("Loaded LoRAs:", lora_names)
         if image_input is not None:
             pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)

flux.py

@@ -17,54 +17,11 @@ from env import (HF_TOKEN, hf_read_token, # to use only for private repos
     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,
+    safe_float, escape_lora_basename, to_lora_key, to_lora_path, get_local_model_list, download_things,
     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, MODEL_TYPE_DICT)
 
 
-def download_things(directory, url, hf_token="", civitai_api_key=""):
-    url = url.strip()
-    if "drive.google.com" in url:
-        original_dir = os.getcwd()
-        os.chdir(directory)
-        os.system(f"gdown --fuzzy {url}")
-        os.chdir(original_dir)
-    elif "huggingface.co" in url:
-        url = url.replace("?download=true", "")
-        # url = urllib.parse.quote(url, safe=':/')  # fix encoding
-        if "/blob/" in url:
-            url = url.replace("/blob/", "/resolve/")
-        user_header = f'"Authorization: Bearer {hf_token}"'
-        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]}")
-    elif "civitai.com" in url:
-        if "?" in url:
-            url = url.split("?")[0]
-        if civitai_api_key:
-            url = url + f"?token={civitai_api_key}"
-            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
-        else:
-            print("\033[91mYou need an API key to download Civitai models.\033[0m")
-    else:
-        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
-
-
-def get_model_list(directory_path):
-    model_list = []
-    valid_extensions = {'.ckpt' , '.pt', '.pth', '.safetensors', '.bin'}
-
-    for filename in os.listdir(directory_path):
-        if os.path.splitext(filename)[1] in valid_extensions:
-            name_without_extension = os.path.splitext(filename)[0]
-            file_path = os.path.join(directory_path, filename)
-            # model_list.append((name_without_extension, file_path))
-            model_list.append(file_path)
-            print('\033[34mFILE: ' + file_path + '\033[0m')
-    return model_list
-
-
 # - **Download Models**
 download_model = ", ".join(download_model_list)
 # - **Download VAEs**
@@ -90,40 +47,10 @@ for url in [url.strip() for url in download_lora.split(',')]:
         download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
 
 lora_model_list = get_lora_model_list()
-vae_model_list = get_model_list(directory_vaes)
+vae_model_list = get_local_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):
-    api = HfApi(token=HF_TOKEN)
-    try:
-        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)
-        return ""
-    if model.private or model.gated: return ""
-    tags = model.tags
-    info = []
-    url = f"https://huggingface.co/{repo_id}/"
-    if not 'diffusers' in tags: return ""
-    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}")
-    info.append(f"likes: {model.likes}")
-    info.append(model.last_modified.strftime("lastmod: %Y-%m-%d"))
-    md = f"Model Info: {', '.join(info)}, [Model Repo]({url})"
-    return gr.update(value=md)
-
-
 private_lora_dict = {"": ["", "", "", "", ""]}
 try:
     with open('lora_dict.json', encoding='utf-8') as f:

modutils.py

@@ -1,10 +1,12 @@
 import spaces
 import json
 import gradio as gr
-from huggingface_hub import HfApi
 import os
 from pathlib import Path
 from PIL import Image
+from huggingface_hub import HfApi, HfFolder, hf_hub_download, snapshot_download
+import urllib.parse
+import re
 
 
 from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
@@ -28,7 +30,7 @@ def to_list(s):
 
 
 def list_uniq(l):
-        return sorted(set(l), key=l.index)
+    return sorted(set(l), key=l.index)
 
 
 def list_sub(a, b):
@@ -61,6 +63,48 @@ def get_local_model_list(dir_path):
     return model_list
 
 
+def get_token():
+    try:
+        token = HfFolder.get_token()
+    except Exception:
+        token = ""
+    return token
+
+
+def set_token(token):
+    try:
+        HfFolder.save_token(token)
+    except Exception:
+        print(f"Error: Failed to save token.")
+
+
+set_token(HF_TOKEN)
+
+
+def split_hf_url(url: str):
+    try:
+        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0])
+        if len(s) < 4: return "", "", "", ""
+        repo_id = s[1]
+        repo_type = "dataset" if s[0] == "datasets" else "model"
+        subfolder = urllib.parse.unquote(s[2]) if s[2] else None
+        filename = urllib.parse.unquote(s[3])
+        return repo_id, filename, subfolder, repo_type
+    except Exception as e:
+        print(e)
+
+
+def download_hf_file(directory, url, progress=gr.Progress(track_tqdm=True)):
+    hf_token = get_token()
+    repo_id, filename, subfolder, repo_type = split_hf_url(url)
+    try:
+        print(f"Downloading {url} to {directory}")
+        if subfolder is not None: hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token)
+        else: hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token)
+    except Exception as e:
+        print(f"Failed to download: {e}")
+
+
 def download_things(directory, url, hf_token="", civitai_api_key=""):
     url = url.strip()
     if "drive.google.com" in url:
@@ -73,11 +117,7 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
         # url = urllib.parse.quote(url, safe=':/')  # fix encoding
         if "/blob/" in url:
             url = url.replace("/blob/", "/resolve/")
-        user_header = f'"Authorization: Bearer {hf_token}"'
-        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]}")
+        download_hf_file(directory, url)
     elif "civitai.com" in url:
         if "?" in url:
             url = url.split("?")[0]
@@ -90,6 +130,33 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
         os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
 
 
+def get_download_file(temp_dir, url, civitai_key="", progress=gr.Progress(track_tqdm=True)):
+    if not "http" in url and is_repo_name(url) and not Path(url).exists():
+        print(f"Use HF Repo: {url}")
+        new_file = url
+    elif not "http" in url and Path(url).exists():
+        print(f"Use local file: {url}")
+        new_file = url
+    elif Path(f"{temp_dir}/{url.split('/')[-1]}").exists():
+        print(f"File to download alreday exists: {url}")
+        new_file = f"{temp_dir}/{url.split('/')[-1]}"
+    else:
+        print(f"Start downloading: {url}")
+        before = get_local_model_list(temp_dir)
+        try:
+            download_things(temp_dir, url.strip(), HF_TOKEN, civitai_key)
+        except Exception:
+            print(f"Download failed: {url}")
+            return ""
+        after = get_local_model_list(temp_dir)
+        new_file = list_sub(after, before)[0] if list_sub(after, before) else ""
+    if not new_file:
+        print(f"Download failed: {url}")
+        return ""
+    print(f"Download completed: {url}")
+    return new_file
+
+
 def escape_lora_basename(basename: str):
     return basename.replace(".", "_").replace(" ", "_").replace(",", "")