lora info

utils.py

@@ -1,562 +1,714 @@
-import os
-import re
-import gradio as gr
-from constants import (
-    DIFFUSERS_FORMAT_LORAS,
-    CIVITAI_API_KEY,
-    HF_TOKEN,
-    MODEL_TYPE_CLASS,
-    DIRECTORY_LORAS,
-    DIRECTORY_MODELS,
-    DIFFUSECRAFT_CHECKPOINT_NAME,
-    CACHE_HF_ROOT,
-    CACHE_HF,
-    STORAGE_ROOT,
-)
-from huggingface_hub import HfApi
-from huggingface_hub import snapshot_download
-from diffusers import DiffusionPipeline
-from huggingface_hub import model_info as model_info_data
-from diffusers.pipelines.pipeline_loading_utils import variant_compatible_siblings
-from stablepy.diffusers_vanilla.utils import checkpoint_model_type
-from pathlib import PosixPath
-from unidecode import unidecode
-import urllib.parse
-import copy
-import requests
-from requests.adapters import HTTPAdapter
-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'
-
-
-def request_json_data(url):
-    model_version_id = url.split('/')[-1]
-    if "?modelVersionId=" in model_version_id:
-        match = re.search(r'modelVersionId=(\d+)', url)
-        model_version_id = match.group(1)
-
-    endpoint_url = f"https://civitai.com/api/v1/model-versions/{model_version_id}"
-
-    params = {}
-    headers = {'User-Agent': USER_AGENT, 'content-type': 'application/json'}
-    session = requests.Session()
-    retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
-    session.mount("https://", HTTPAdapter(max_retries=retries))
-
-    try:
-        result = session.get(endpoint_url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
-        result.raise_for_status()
-        json_data = result.json()
-        return json_data if json_data else None
-    except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
-class ModelInformation:
-    def __init__(self, json_data):
-        self.model_version_id = json_data.get("id", "")
-        self.model_id = json_data.get("modelId", "")
-        self.download_url = json_data.get("downloadUrl", "")
-        self.model_url = f"https://civitai.com/models/{self.model_id}?modelVersionId={self.model_version_id}"
-        self.filename_url = next(
-            (v.get("name", "") for v in json_data.get("files", []) if str(self.model_version_id) in v.get("downloadUrl", "") and v.get("type", "Model") == "Model"), ""
-        )
-        self.filename_url = self.filename_url if self.filename_url else ""
-        self.description = json_data.get("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)
-        self.poi = json_data.get("model", {}).get("poi", False)
-        self.images = [img.get("url", "") for img in json_data.get("images", [])]
-        self.example_prompt = json_data.get("trainedWords", [""])[0] if json_data.get("trainedWords") else ""
-        self.original_json = copy.deepcopy(json_data)
-
-
-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 civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
-    filename_base = filename = None
-
-    if alternative_name:
-        output_path = os.path.join(dir_, alternative_name)
-        if os.path.exists(output_path):
-            return output_path, alternative_name
-
-    # 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 redirect: {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 "{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
-
-    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}")
-
-    return downloaded_file_path
-
-
-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
-
-
-def extract_parameters(input_string):
-    parameters = {}
-    input_string = input_string.replace("\n", "")
-
-    if "Negative prompt:" not in input_string:
-        if "Steps:" in input_string:
-            input_string = input_string.replace("Steps:", "Negative prompt: Steps:")
-        else:
-            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]:
-        parameters["neg_prompt"] = parm[1].strip()
-        return parameters
-    parm = parm[1].split("Steps:")
-    parameters["neg_prompt"] = parm[0].strip()
-    input_string = "Steps:" + parm[1]
-
-    # Extracting Steps
-    steps_match = re.search(r'Steps: (\d+)', input_string)
-    if steps_match:
-        parameters['Steps'] = int(steps_match.group(1))
-
-    # Extracting Size
-    size_match = re.search(r'Size: (\d+x\d+)', input_string)
-    if size_match:
-        parameters['Size'] = size_match.group(1)
-        width, height = map(int, parameters['Size'].split('x'))
-        parameters['width'] = width
-        parameters['height'] = height
-
-    # Extracting other parameters
-    other_parameters = re.findall(r'([^,:]+): (.*?)(?=, [^,:]+:|$)', input_string)
-    for param in other_parameters:
-        parameters[param[0].strip()] = param[1].strip('"')
-
-    return parameters
-
-
-def get_my_lora(link_url, romanize):
-    l_name = ""
-    for url in [url.strip() for url in link_url.split(',')]:
-        if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
-            l_name = download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY, romanize)
-    new_lora_model_list = get_model_list(DIRECTORY_LORAS)
-    new_lora_model_list.insert(0, "None")
-    new_lora_model_list = new_lora_model_list + DIFFUSERS_FORMAT_LORAS
-    msg_lora = "Downloaded"
-    if l_name:
-        msg_lora += f": <b>{l_name}</b>"
-        print(msg_lora)
-
-    return gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        choices=new_lora_model_list
-    ), gr.update(
-        value=msg_lora
-    )
-
-
-def info_html(json_data, title, subtitle):
-    return f"""
-        <div style='padding: 0; border-radius: 10px;'>
-            <p style='margin: 0; font-weight: bold;'>{title}</p>
-            <details>
-                <summary>Details</summary>
-                <p style='margin: 0; font-weight: bold;'>{subtitle}</p>
-            </details>
-        </div>
-        """
-
-
-def get_model_type(repo_id: str):
-    api = HfApi(token=os.environ.get("HF_TOKEN"))  # if use private or gated model
-    default = "SD 1.5"
-    try:
-        if os.path.exists(repo_id):
-            tag, _, _, _ = checkpoint_model_type(repo_id)
-            return DIFFUSECRAFT_CHECKPOINT_NAME[tag]
-        else:
-            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)
-
-    except Exception:
-        return default
-    return default
-
-
-def restart_space(repo_id: str, factory_reboot: bool):
-    api = HfApi(token=os.environ.get("HF_TOKEN"))
-    try:
-        runtime = api.get_space_runtime(repo_id=repo_id)
-        if runtime.stage == "RUNNING":
-            api.restart_space(repo_id=repo_id, factory_reboot=factory_reboot)
-            print(f"Restarting space: {repo_id}")
-        else:
-            print(f"Space {repo_id} is in stage: {runtime.stage}")
-    except Exception as e:
-        print(e)
-
-
-def extract_exif_data(image):
-    if image is None:
-        return ""
-
-    try:
-        metadata_keys = ['parameters', 'metadata', 'prompt', 'Comment']
-
-        for key in metadata_keys:
-            if key in image.info:
-                return image.info[key]
-
-        return str(image.info)
-
-    except Exception as e:
-        return f"Error extracting metadata: {str(e)}"
-
-
-def create_mask_now(img, invert):
-    import numpy as np
-    import time
-
-    time.sleep(0.5)
-
-    transparent_image = img["layers"][0]
-
-    # Extract the alpha channel
-    alpha_channel = np.array(transparent_image)[:, :, 3]
-
-    # Create a binary mask by thresholding the alpha channel
-    binary_mask = alpha_channel > 1
-
-    if invert:
-        print("Invert")
-        # Invert the binary mask so that the drawn shape is white and the rest is black
-        binary_mask = np.invert(binary_mask)
-
-    # Convert the binary mask to a 3-channel RGB mask
-    rgb_mask = np.stack((binary_mask,) * 3, axis=-1)
-
-    # Convert the mask to uint8
-    rgb_mask = rgb_mask.astype(np.uint8) * 255
-
-    return img["background"], rgb_mask
-
-
-def download_diffuser_repo(repo_name: str, model_type: str, revision: str = "main", token=True):
-
-    variant = None
-    if token is True and not os.environ.get("HF_TOKEN"):
-        token = None
-
-    if model_type == "SDXL":
-        info = model_info_data(
-            repo_name,
-            token=token,
-            revision=revision,
-            timeout=5.0,
-        )
-
-        filenames = {sibling.rfilename for sibling in info.siblings}
-        model_filenames, variant_filenames = variant_compatible_siblings(
-            filenames, variant="fp16"
-        )
-
-        if len(variant_filenames):
-            variant = "fp16"
-
-    if model_type == "FLUX":
-        cached_folder = snapshot_download(
-            repo_id=repo_name,
-            allow_patterns="transformer/*"
-        )
-    else:
-        cached_folder = DiffusionPipeline.download(
-            pretrained_model_name=repo_name,
-            force_download=False,
-            token=token,
-            revision=revision,
-            # mirror="https://hf-mirror.com",
-            variant=variant,
-            use_safetensors=True,
-            trust_remote_code=False,
-            timeout=5.0,
-        )
-
-    if isinstance(cached_folder, PosixPath):
-        cached_folder = cached_folder.as_posix()
-
-    # Task model
-    # from huggingface_hub import hf_hub_download
-    # hf_hub_download(
-    #     task_model,
-    #     filename="diffusion_pytorch_model.safetensors",  # fix fp16 variant
-    # )
-
-    return cached_folder
-
-
-def get_folder_size_gb(folder_path):
-    result = subprocess.run(["du", "-s", folder_path], capture_output=True, text=True)
-
-    total_size_kb = int(result.stdout.split()[0])
-    total_size_gb = total_size_kb / (1024 ** 2)
-
-    return total_size_gb
-
-
-def get_used_storage_gb(path_storage=STORAGE_ROOT):
-    try:
-        used_gb = get_folder_size_gb(path_storage)
-        print(f"Used Storage: {used_gb:.2f} GB")
-    except Exception as e:
-        used_gb = 999
-        print(f"Error while retrieving the used storage: {e}.")
-
-    return used_gb
-
-
-def delete_model(removal_candidate):
-    print(f"Removing: {removal_candidate}")
-
-    if os.path.exists(removal_candidate):
-        os.remove(removal_candidate)
-    else:
-        diffusers_model = f"{CACHE_HF}{DIRECTORY_MODELS}--{removal_candidate.replace('/', '--')}"
-        if os.path.isdir(diffusers_model):
-            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))
-
-    return f"""
-        <div style="position: relative; width: 100%; background-color: gray; border-radius: 5px; overflow: hidden;">
-            <div style="width: {percentage}%; height: 17px; background-color: #800080; transition: width 0.5s;"></div>
-            <div style="position: absolute; width: 100%; text-align: center; color: white; top: 0; line-height: 19px; font-size: 13px;">
-                {int(percentage)}%
-            </div>
-        </div>
-        """
-
-
-def html_template_message(msg):
-    return f"""
-        <div style="position: relative; width: 100%; background-color: gray; border-radius: 5px; overflow: hidden;">
-            <div style="width: 0%; height: 17px; background-color: #800080; transition: width 0.5s;"></div>
-            <div style="position: absolute; width: 100%; text-align: center; color: white; top: 0; line-height: 19px; font-size: 14px; font-weight: bold; text-shadow: 1px 1px 2px black;">
-                {msg}
-            </div>
-        </div>
-        """
-
-
-def escape_html(text):
-    """Escapes HTML special characters in the input text."""
-    return text.replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")
+import os
+import re
+import gradio as gr
+from constants import (
+    DIFFUSERS_FORMAT_LORAS,
+    CIVITAI_API_KEY,
+    HF_TOKEN,
+    MODEL_TYPE_CLASS,
+    DIRECTORY_LORAS,
+    DIRECTORY_MODELS,
+    DIFFUSECRAFT_CHECKPOINT_NAME,
+    CACHE_HF_ROOT,
+    CACHE_HF,
+    STORAGE_ROOT,
+)
+from huggingface_hub import HfApi, get_hf_file_metadata, snapshot_download
+from diffusers import DiffusionPipeline
+from huggingface_hub import model_info as model_info_data
+from diffusers.pipelines.pipeline_loading_utils import variant_compatible_siblings
+from stablepy.diffusers_vanilla.utils import checkpoint_model_type
+from pathlib import PosixPath
+from unidecode import unidecode
+import urllib.parse
+import copy
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
+import shutil
+import subprocess
+import json
+import html as _html
+
+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'
+MODEL_ARCH = {
+    'stable-diffusion-xl-v1-base/lora': "Stable Diffusion XL (Illustrious, Pony, NoobAI)",
+    'stable-diffusion-v1/lora': "Stable Diffusion 1.5",
+    'flux-1-dev/lora': "Flux",
+}
+
+
+def read_safetensors_header_from_url(url: str):
+    """Read safetensors header from a remote Hugging Face file."""
+    meta = get_hf_file_metadata(url)
+
+    # Step 1: first 8 bytes → header length
+    resp = requests.get(meta.location, headers={"Range": "bytes=0-7"})
+    resp.raise_for_status()
+    header_len = int.from_bytes(resp.content, "little")
+
+    # Step 2: fetch full header JSON
+    end = 8 + header_len - 1
+    resp = requests.get(meta.location, headers={"Range": f"bytes=8-{end}"})
+    resp.raise_for_status()
+    header_json = resp.content.decode("utf-8")
+
+    return json.loads(header_json)
+
+
+def read_safetensors_header_from_file(path: str):
+    """Read safetensors header from a local file."""
+    with open(path, "rb") as f:
+        # Step 1: first 8 bytes → header length
+        header_len = int.from_bytes(f.read(8), "little")
+
+        # Step 2: read header JSON
+        header_json = f.read(header_len).decode("utf-8")
+
+    return json.loads(header_json)
+
+
+class LoraHeaderInformation:
+    """
+    Encapsulates parsed info from a LoRA JSON header and provides
+    a compact HTML summary via .to_html().
+    """
+
+    def __init__(self, json_data):
+        self.original_json = copy.deepcopy(json_data or {})
+
+        # Check if text encoder was trained
+        # guard for json_data being a mapping
+        try:
+            self.text_encoder_trained = any("text_model" in ln for ln in json_data)
+        except Exception:
+            self.text_encoder_trained = False
+
+        # Metadata (may be None)
+        metadata = (json_data or {}).get("__metadata__", None)
+        self.metadata = metadata
+
+        # Default values
+        self.architecture = "undefined"
+        self.prediction_type = "undefined"
+        self.base_model = "undefined"
+        self.author = "undefined"
+        self.title = "undefined"
+        self.common_tags_list = []
+
+        if metadata:
+            self.architecture = MODEL_ARCH.get(
+                metadata.get('modelspec.architecture', None),
+                "undefined"
+            )
+
+            self.prediction_type = metadata.get('modelspec.prediction_type', "undefined")
+            self.base_model = metadata.get('ss_sd_model_name', "undefined")
+            self.author = metadata.get('modelspec.author', "undefined")
+            self.title = metadata.get('modelspec.title', "undefined")
+
+            base_model_hash = metadata.get('ss_new_sd_model_hash', None)  # SHA256
+            # AUTOV1 ss_sd_model_hash
+            # https://civitai.com/api/v1/model-versions/by-hash/{base_model_hash}  # Info
+            if base_model_hash:
+                self.base_model += f" hash={base_model_hash}"
+
+            # Extract tags
+            try:
+                tags = metadata.get('ss_tag_frequency') if "ss_tag_frequency" in metadata else metadata.get('ss_datasets', "")
+                tags = json.loads(tags) if tags else ""
+
+                if isinstance(tags, list):
+                    tags = tags[0].get("tag_frequency", {})
+
+                if tags:
+                    self.common_tags_list = list(tags[list(tags.keys())[0]].keys())
+            except Exception:
+                self.common_tags_list = []
+
+    def to_dict(self):
+        """Return a plain dict summary of parsed fields."""
+        return {
+            "architecture": self.architecture,
+            "prediction_type": self.prediction_type,
+            "base_model": self.base_model,
+            "author": self.author,
+            "title": self.title,
+            "text_encoder_trained": bool(self.text_encoder_trained),
+            "common_tags": self.common_tags_list,
+        }
+
+    def to_html(self, limit_tags=20):
+        """
+        Return a compact HTML snippet (string) showing the parsed info
+        in a small font. Values are HTML-escaped.
+        """
+        # helper to escape
+        esc = _html.escape
+
+        rows = [
+            ("Title", esc(str(self.title))),
+            ("Author", esc(str(self.author))),
+            ("Architecture", esc(str(self.architecture))),
+            ("Base model", esc(str(self.base_model))),
+            ("Prediction type", esc(str(self.prediction_type))),
+            ("Text encoder trained", esc(str(self.text_encoder_trained))),
+            ("Reference tags", esc(str(", ".join(self.common_tags_list[:limit_tags])))),
+        ]
+
+        # small, compact table with inline styling (small font)
+        html_rows = "".join(
+            f"<tr><th style='text-align:left;padding:2px 6px;white-space:nowrap'>{k}</th>"
+            f"<td style='padding:2px 6px'>{v}</td></tr>"
+            for k, v in rows
+        )
+
+        html_snippet = (
+            "<div style='font-family:system-ui, -apple-system, \"Segoe UI\", Roboto, "
+            "Helvetica, Arial, \"Noto Sans\", sans-serif; font-size:12px; line-height:1.2; "
+            "'>"
+            f"<table style='border-collapse:collapse; font-size:12px;'>"
+            f"{html_rows}"
+            "</table>"
+            "</div>"
+        )
+
+        return html_snippet
+
+
+def request_json_data(url):
+    model_version_id = url.split('/')[-1]
+    if "?modelVersionId=" in model_version_id:
+        match = re.search(r'modelVersionId=(\d+)', url)
+        model_version_id = match.group(1)
+
+    endpoint_url = f"https://civitai.com/api/v1/model-versions/{model_version_id}"
+
+    params = {}
+    headers = {'User-Agent': USER_AGENT, 'content-type': 'application/json'}
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
+    session.mount("https://", HTTPAdapter(max_retries=retries))
+
+    try:
+        result = session.get(endpoint_url, params=params, headers=headers, stream=True, timeout=(3.0, 15))
+        result.raise_for_status()
+        json_data = result.json()
+        return json_data if json_data else None
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+class ModelInformation:
+    def __init__(self, json_data):
+        self.model_version_id = json_data.get("id", "")
+        self.model_id = json_data.get("modelId", "")
+        self.download_url = json_data.get("downloadUrl", "")
+        self.model_url = f"https://civitai.com/models/{self.model_id}?modelVersionId={self.model_version_id}"
+        self.filename_url = next(
+            (v.get("name", "") for v in json_data.get("files", []) if str(self.model_version_id) in v.get("downloadUrl", "") and v.get("type", "Model") == "Model"), ""
+        )
+        self.filename_url = self.filename_url if self.filename_url else ""
+        self.description = json_data.get("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)
+        self.poi = json_data.get("model", {}).get("poi", False)
+        self.images = [img.get("url", "") for img in json_data.get("images", [])]
+        self.example_prompt = json_data.get("trainedWords", [""])[0] if json_data.get("trainedWords") else ""
+        self.original_json = copy.deepcopy(json_data)
+
+
+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 civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
+    filename_base = filename = None
+
+    if alternative_name:
+        output_path = os.path.join(dir_, alternative_name)
+        if os.path.exists(output_path):
+            return output_path, alternative_name
+
+    # 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 redirect: {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 "{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
+
+    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}")
+
+    return downloaded_file_path
+
+
+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
+
+
+def extract_parameters(input_string):
+    parameters = {}
+    input_string = input_string.replace("\n", "")
+
+    if "Negative prompt:" not in input_string:
+        if "Steps:" in input_string:
+            input_string = input_string.replace("Steps:", "Negative prompt: Steps:")
+        else:
+            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]:
+        parameters["neg_prompt"] = parm[1].strip()
+        return parameters
+    parm = parm[1].split("Steps:")
+    parameters["neg_prompt"] = parm[0].strip()
+    input_string = "Steps:" + parm[1]
+
+    # Extracting Steps
+    steps_match = re.search(r'Steps: (\d+)', input_string)
+    if steps_match:
+        parameters['Steps'] = int(steps_match.group(1))
+
+    # Extracting Size
+    size_match = re.search(r'Size: (\d+x\d+)', input_string)
+    if size_match:
+        parameters['Size'] = size_match.group(1)
+        width, height = map(int, parameters['Size'].split('x'))
+        parameters['width'] = width
+        parameters['height'] = height
+
+    # Extracting other parameters
+    other_parameters = re.findall(r'([^,:]+): (.*?)(?=, [^,:]+:|$)', input_string)
+    for param in other_parameters:
+        parameters[param[0].strip()] = param[1].strip('"')
+
+    return parameters
+
+
+def get_my_lora(link_url, romanize):
+    l_name = ""
+    for url in [url.strip() for url in link_url.split(',')]:
+        if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
+            l_name = download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY, romanize)
+    new_lora_model_list = get_model_list(DIRECTORY_LORAS)
+    new_lora_model_list.insert(0, "None")
+    new_lora_model_list = new_lora_model_list + DIFFUSERS_FORMAT_LORAS
+    msg_lora = "Downloaded"
+    if l_name:
+        msg_lora += f": <b>{l_name}</b>"
+        print(msg_lora)
+
+    try:
+        # Works with non-Civitai loras.
+        json_data = read_safetensors_header_from_file(l_name)
+        metadata_lora = LoraHeaderInformation(json_data)
+        msg_lora += "<br>" + metadata_lora.to_html()
+    except Exception:
+        pass
+
+    return gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        choices=new_lora_model_list
+    ), gr.update(
+        value=msg_lora
+    )
+
+
+def info_html(json_data, title, subtitle):
+    return f"""
+        <div style='padding: 0; border-radius: 10px;'>
+            <p style='margin: 0; font-weight: bold;'>{title}</p>
+            <details>
+                <summary>Details</summary>
+                <p style='margin: 0; font-weight: bold;'>{subtitle}</p>
+            </details>
+        </div>
+        """
+
+
+def get_model_type(repo_id: str):
+    api = HfApi(token=os.environ.get("HF_TOKEN"))  # if use private or gated model
+    default = "SD 1.5"
+    try:
+        if os.path.exists(repo_id):
+            tag, _, _, _ = checkpoint_model_type(repo_id)
+            return DIFFUSECRAFT_CHECKPOINT_NAME[tag]
+        else:
+            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)
+
+    except Exception:
+        return default
+    return default
+
+
+def restart_space(repo_id: str, factory_reboot: bool):
+    api = HfApi(token=os.environ.get("HF_TOKEN"))
+    try:
+        runtime = api.get_space_runtime(repo_id=repo_id)
+        if runtime.stage == "RUNNING":
+            api.restart_space(repo_id=repo_id, factory_reboot=factory_reboot)
+            print(f"Restarting space: {repo_id}")
+        else:
+            print(f"Space {repo_id} is in stage: {runtime.stage}")
+    except Exception as e:
+        print(e)
+
+
+def extract_exif_data(image):
+    if image is None:
+        return ""
+
+    try:
+        metadata_keys = ['parameters', 'metadata', 'prompt', 'Comment']
+
+        for key in metadata_keys:
+            if key in image.info:
+                return image.info[key]
+
+        return str(image.info)
+
+    except Exception as e:
+        return f"Error extracting metadata: {str(e)}"
+
+
+def create_mask_now(img, invert):
+    import numpy as np
+    import time
+
+    time.sleep(0.5)
+
+    transparent_image = img["layers"][0]
+
+    # Extract the alpha channel
+    alpha_channel = np.array(transparent_image)[:, :, 3]
+
+    # Create a binary mask by thresholding the alpha channel
+    binary_mask = alpha_channel > 1
+
+    if invert:
+        print("Invert")
+        # Invert the binary mask so that the drawn shape is white and the rest is black
+        binary_mask = np.invert(binary_mask)
+
+    # Convert the binary mask to a 3-channel RGB mask
+    rgb_mask = np.stack((binary_mask,) * 3, axis=-1)
+
+    # Convert the mask to uint8
+    rgb_mask = rgb_mask.astype(np.uint8) * 255
+
+    return img["background"], rgb_mask
+
+
+def download_diffuser_repo(repo_name: str, model_type: str, revision: str = "main", token=True):
+
+    variant = None
+    if token is True and not os.environ.get("HF_TOKEN"):
+        token = None
+
+    if model_type == "SDXL":
+        info = model_info_data(
+            repo_name,
+            token=token,
+            revision=revision,
+            timeout=5.0,
+        )
+
+        filenames = {sibling.rfilename for sibling in info.siblings}
+        model_filenames, variant_filenames = variant_compatible_siblings(
+            filenames, variant="fp16"
+        )
+
+        if len(variant_filenames):
+            variant = "fp16"
+
+    if model_type == "FLUX":
+        cached_folder = snapshot_download(
+            repo_id=repo_name,
+            allow_patterns="transformer/*"
+        )
+    else:
+        cached_folder = DiffusionPipeline.download(
+            pretrained_model_name=repo_name,
+            force_download=False,
+            token=token,
+            revision=revision,
+            # mirror="https://hf-mirror.com",
+            variant=variant,
+            use_safetensors=True,
+            trust_remote_code=False,
+            timeout=5.0,
+        )
+
+    if isinstance(cached_folder, PosixPath):
+        cached_folder = cached_folder.as_posix()
+
+    # Task model
+    # from huggingface_hub import hf_hub_download
+    # hf_hub_download(
+    #     task_model,
+    #     filename="diffusion_pytorch_model.safetensors",  # fix fp16 variant
+    # )
+
+    return cached_folder
+
+
+def get_folder_size_gb(folder_path):
+    result = subprocess.run(["du", "-s", folder_path], capture_output=True, text=True)
+
+    total_size_kb = int(result.stdout.split()[0])
+    total_size_gb = total_size_kb / (1024 ** 2)
+
+    return total_size_gb
+
+
+def get_used_storage_gb(path_storage=STORAGE_ROOT):
+    try:
+        used_gb = get_folder_size_gb(path_storage)
+        print(f"Used Storage: {used_gb:.2f} GB")
+    except Exception as e:
+        used_gb = 999
+        print(f"Error while retrieving the used storage: {e}.")
+
+    return used_gb
+
+
+def delete_model(removal_candidate):
+    print(f"Removing: {removal_candidate}")
+
+    if os.path.exists(removal_candidate):
+        os.remove(removal_candidate)
+    else:
+        diffusers_model = f"{CACHE_HF}{DIRECTORY_MODELS}--{removal_candidate.replace('/', '--')}"
+        if os.path.isdir(diffusers_model):
+            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):
+            shutil.rmtree(CACHE_HF, ignore_errors=True)
+            print(f"Hugging Face cache cleared: {CACHE_HF}")
+        else:
+            print(f"No Hugging Face cache found at: {CACHE_HF}")
+    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))
+
+    return f"""
+        <div style="position: relative; width: 100%; background-color: gray; border-radius: 5px; overflow: hidden;">
+            <div style="width: {percentage}%; height: 17px; background-color: #800080; transition: width 0.5s;"></div>
+            <div style="position: absolute; width: 100%; text-align: center; color: white; top: 0; line-height: 19px; font-size: 13px;">
+                {int(percentage)}%
+            </div>
+        </div>
+        """
+
+
+def html_template_message(msg):
+    return f"""
+        <div style="position: relative; width: 100%; background-color: gray; border-radius: 5px; overflow: hidden;">
+            <div style="width: 0%; height: 17px; background-color: #800080; transition: width 0.5s;"></div>
+            <div style="position: absolute; width: 100%; text-align: center; color: white; top: 0; line-height: 19px; font-size: 14px; font-weight: bold; text-shadow: 1px 1px 2px black;">
+                {msg}
+            </div>
+        </div>
+        """
+
+
+def escape_html(text):
+    """Escapes HTML special characters in the input text."""
+    return text.replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")