Upload utils.py

utils.py

@@ -9,6 +9,7 @@ from constants import (
     DIRECTORY_LORAS,
     DIRECTORY_MODELS,
     DIFFUSECRAFT_CHECKPOINT_NAME,
+    CACHE_HF_ROOT,
     CACHE_HF,
     STORAGE_ROOT,
 )
@@ -28,6 +29,7 @@ from urllib3.util import Retry
 import shutil
 import subprocess
 
+IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
 USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'
 
 
@@ -66,7 +68,8 @@ class ModelInformation:
         )
         self.filename_url = self.filename_url if self.filename_url else ""
         self.description = json_data.get("description", "")
-        if self.description is None: self.description = ""
+        if self.description is None:
+            self.description = ""
         self.model_name = json_data.get("model", {}).get("name", "")
         self.model_type = json_data.get("model", {}).get("type", "")
         self.nsfw = json_data.get("model", {}).get("nsfw", False)
@@ -76,118 +79,175 @@ class ModelInformation:
         self.original_json = copy.deepcopy(json_data)
 
 
-def retrieve_model_info(url):
-    json_data = request_json_data(url)
-    if not json_data:
-        return None
-    model_descriptor = ModelInformation(json_data)
-    return model_descriptor
+def get_civit_params(url):
+    try:
+        json_data = request_json_data(url)
+        mdc = ModelInformation(json_data)
+        if mdc.download_url and mdc.filename_url:
+            return mdc.download_url, mdc.filename_url, mdc.model_url
+        else:
+            ValueError("Invalid Civitai model URL")
+    except Exception as e:
+        print(f"Error retrieving Civitai metadata: {e} — fallback to direct download")
+        return url, None, None
 
 
-def download_things(directory, url, hf_token="", civitai_api_key="", romanize=False):
-    url = url.strip()
-    downloaded_file_path = None
+def civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
+    filename_base = filename = None
 
-    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 alternative_name:
+        output_path = os.path.join(dir_, alternative_name)
+        if os.path.exists(output_path):
+            return output_path, alternative_name
 
-        filename = unidecode(url.split('/')[-1]) if romanize else url.split('/')[-1]
+    # Follow the redirect to get the actual download URL
+    curl_command = (
+        f'curl -L -sI --connect-timeout 5 --max-time 5 '
+        f'-H "Content-Type: application/json" '
+        f'-H "Authorization: Bearer {civitai_api_key}" "{url}"'
+    )
 
-        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 {filename}")
-        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 {filename}")
+    headers = os.popen(curl_command).read()
 
-        downloaded_file_path = os.path.join(directory, filename)
+    # Look for the redirected "Location" URL
+    location_match = re.search(r'location: (.+)', headers, re.IGNORECASE)
 
-    elif "civitai.com" in url:
+    if location_match:
+        redirect_url = location_match.group(1).strip()
 
-        if not civitai_api_key:
-            print("\033[91mYou need an API key to download Civitai models.\033[0m")
-
-        model_profile = retrieve_model_info(url)
-        if (
-            model_profile is not None
-            and model_profile.download_url
-            and model_profile.filename_url
-        ):
-            url = model_profile.download_url
-            filename = unidecode(model_profile.filename_url) if romanize else model_profile.filename_url
-        else:
-            if "?" in url:
-                url = url.split("?")[0]
-            filename = ""
+        # Extract the filename from the redirect URL's "Content-Disposition"
+        filename_match = re.search(r'filename%3D%22(.+?)%22', redirect_url)
+        if filename_match:
+            encoded_filename = filename_match.group(1)
+            # Decode the URL-encoded filename
+            decoded_filename = urllib.parse.unquote(encoded_filename)
 
-        url_dl = url + f"?token={civitai_api_key}"
-        print(f"Filename: {filename}")
+            filename = unidecode(decoded_filename) if romanize else decoded_filename
+            # print(f"Filename redirect: {filename}")
 
-        param_filename = ""
-        if filename:
-            param_filename = f"-o '{filename}'"
+    filename_base = alternative_name if alternative_name else filename
+    if not filename_base:
+        return None, None
+    elif os.path.exists(os.path.join(dir_, filename_base)):
+        return os.path.join(dir_, filename_base), filename_base
 
-        aria2_command = (
-            f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
-            f'-k 1M -s 16 -d "{directory}" {param_filename} "{url_dl}"'
-        )
+    aria2_command = (
+        f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
+        f'-k 1M -s 16 -d "{dir_}" -o "{filename_base}" "{redirect_url}"'
+    )
+    r_code = os.system(aria2_command)  # noqa
+
+    # if r_code != 0:
+    #     raise RuntimeError(f"Failed to download file: {filename_base}. Error code: {r_code}")
+
+    output_path = os.path.join(dir_, filename_base)
+    if not os.path.exists(output_path):
+        return None, filename_base
+
+    return output_path, filename_base
+
+
+def civ_api_down(url, dir_, civitai_api_key, civ_filename):
+    """
+    This method is susceptible to being blocked because it generates a lot of temp redirect links with aria2c.
+    If an API key limit is reached, generating a new API key and using it can fix the issue.
+    """
+    output_path = None
+
+    url_dl = url + f"?token={civitai_api_key}"
+    if not civ_filename:
+        aria2_command = f'aria2c -c -x 1 -s 1 -d "{dir_}" "{url_dl}"'
         os.system(aria2_command)
+    else:
+        output_path = os.path.join(dir_, civ_filename)
+        if not os.path.exists(output_path):
+            aria2_command = (
+                f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
+                f'-k 1M -s 16 -d "{dir_}" -o "{civ_filename}" "{url_dl}"'
+            )
+            os.system(aria2_command)
+
+    return output_path
+
+
+def drive_down(url, dir_):
+    import gdown
+
+    output_path = None
 
-        if param_filename and os.path.exists(os.path.join(directory, filename)):
-            downloaded_file_path = os.path.join(directory, filename)
-
-        # # PLAN B
-        # # Follow the redirect to get the actual download URL
-        # curl_command = (
-        #     f'curl -L -sI --connect-timeout 5 --max-time 5 '
-        #     f'-H "Content-Type: application/json" '
-        #     f'-H "Authorization: Bearer {civitai_api_key}" "{url}"'
-        # )
-
-        # headers = os.popen(curl_command).read()
-
-        # # Look for the redirected "Location" URL
-        # location_match = re.search(r'location: (.+)', headers, re.IGNORECASE)
-
-        # if location_match:
-        #     redirect_url = location_match.group(1).strip()
-
-        #     # Extract the filename from the redirect URL's "Content-Disposition"
-        #     filename_match = re.search(r'filename%3D%22(.+?)%22', redirect_url)
-        #     if filename_match:
-        #         encoded_filename = filename_match.group(1)
-        #         # Decode the URL-encoded filename
-        #         decoded_filename = urllib.parse.unquote(encoded_filename)
-
-        #         filename = unidecode(decoded_filename) if romanize else decoded_filename
-        #         print(f"Filename: {filename}")
-
-        #         aria2_command = (
-        #             f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
-        #             f'-k 1M -s 16 -d "{directory}" -o "{filename}" "{redirect_url}"'
-        #         )
-        #         return_code = os.system(aria2_command)
-
-        #         # if return_code != 0:
-        #         #     raise RuntimeError(f"Failed to download file: {filename}. Error code: {return_code}")
-        #         downloaded_file_path = os.path.join(directory, filename)
-        #         if not os.path.exists(downloaded_file_path):
-        #             downloaded_file_path = None
-
-        # if not downloaded_file_path:
-        #     # Old method
-        #     if "?" in url:
-        #         url = url.split("?")[0]
-        #     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}")
+    drive_id, _ = gdown.parse_url.parse_url(url, warning=False)
+    dir_files = os.listdir(dir_)
 
+    for dfile in dir_files:
+        if drive_id in dfile:
+            output_path = os.path.join(dir_, dfile)
+            break
+
+    if not output_path:
+        original_path = gdown.download(url, f"{dir_}/", fuzzy=True)
+
+        dir_name, base_name = os.path.split(original_path)
+        name, ext = base_name.rsplit(".", 1)
+        new_name = f"{name}_{drive_id}.{ext}"
+        output_path = os.path.join(dir_name, new_name)
+
+        os.rename(original_path, output_path)
+
+    return output_path
+
+
+def hf_down(url, dir_, hf_token, romanize):
+    url = url.replace("?download=true", "")
+    # url = urllib.parse.quote(url, safe=':/')  # fix encoding
+
+    filename = unidecode(url.split('/')[-1]) if romanize else url.split('/')[-1]
+    output_path = os.path.join(dir_, filename)
+
+    if os.path.exists(output_path):
+        return output_path
+
+    if "/blob/" in url:
+        url = url.replace("/blob/", "/resolve/")
+
+    if hf_token:
+        user_header = f'"Authorization: Bearer {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 {dir_}  -o {filename}")
+    else:
+        os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {dir_}  -o {filename}")
+
+    return output_path
+
+
+def download_things(directory, url, hf_token="", civitai_api_key="", romanize=False):
+    url = url.strip()
+    downloaded_file_path = None
+
+    if "drive.google.com" in url:
+        downloaded_file_path = drive_down(url, directory)
+    elif "huggingface.co" in url:
+        downloaded_file_path = hf_down(url, directory, hf_token, romanize)
+    elif "civitai.com" in url:
+        if not civitai_api_key:
+            msg = "You need an API key to download Civitai models."
+            print(f"\033[91m{msg}\033[0m")
+            gr.Warning(msg)
+            return None
+
+        url, civ_filename, civ_page = get_civit_params(url)
+        if civ_page and not IS_ZERO_GPU:
+            print(f"\033[92mCivitai model: {civ_filename} [page: {civ_page}]\033[0m")
+
+        downloaded_file_path, civ_filename = civ_redirect_down(url, directory, civitai_api_key, romanize, civ_filename)
+
+        if not downloaded_file_path:
+            msg = (
+                "Download failed.\n"
+                "If this is due to an API limit, generating a new API key may resolve the issue.\n"
+                "Attempting to download using the old method..."
+            )
+            print(msg)
+            gr.Warning(msg)
+            downloaded_file_path = civ_api_down(url, directory, civitai_api_key, civ_filename)
     else:
         os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
 
@@ -216,14 +276,15 @@ def extract_parameters(input_string):
         if "Steps:" in input_string:
             input_string = input_string.replace("Steps:", "Negative prompt: Steps:")
         else:
-            print("Invalid metadata")
+            msg = "Generation data is invalid."
+            gr.Warning(msg)
+            print(msg)
             parameters["prompt"] = input_string
             return parameters
 
     parm = input_string.split("Negative prompt:")
     parameters["prompt"] = parm[0].strip()
     if "Steps:" not in parm[1]:
-        print("Steps not detected")
         parameters["neg_prompt"] = parm[1].strip()
         return parameters
     parm = parm[1].split("Steps:")
@@ -306,7 +367,8 @@ def get_model_type(repo_id: str):
             model = api.model_info(repo_id=repo_id, timeout=5.0)
             tags = model.tags
             for tag in tags:
-                if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
+                if tag in MODEL_TYPE_CLASS.keys():
+                    return MODEL_TYPE_CLASS.get(tag, default)
 
     except Exception:
         return default
@@ -433,9 +495,9 @@ def get_folder_size_gb(folder_path):
     return total_size_gb
 
 
-def get_used_storage_gb():
+def get_used_storage_gb(path_storage=STORAGE_ROOT):
     try:
-        used_gb = get_folder_size_gb(STORAGE_ROOT)
+        used_gb = get_folder_size_gb(path_storage)
         print(f"Used Storage: {used_gb:.2f} GB")
     except Exception as e:
         used_gb = 999
@@ -455,6 +517,21 @@ def delete_model(removal_candidate):
             shutil.rmtree(diffusers_model)
 
 
+def clear_hf_cache():
+    """
+    Clears the entire Hugging Face cache at ~/.cache/huggingface.
+    Hugging Face will re-download models as needed later.
+    """
+    try:
+        if os.path.exists(CACHE_HF_ROOT):
+            shutil.rmtree(CACHE_HF_ROOT, ignore_errors=True)
+            print(f"Hugging Face cache cleared: {CACHE_HF_ROOT}")
+        else:
+            print(f"No Hugging Face cache found at: {CACHE_HF_ROOT}")
+    except Exception as e:
+        print(f"Error clearing Hugging Face cache: {e}")
+
+
 def progress_step_bar(step, total):
     # Calculate the percentage for the progress bar width
     percentage = min(100, ((step / total) * 100))