utils.py

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>")