mz_kolors_utils.py

import json
import os
import shutil
import subprocess
import sys
import threading
import time
import numpy as np
import folder_paths
import base64
from PIL import Image, ImageFilter
import io
import torch
import re
import hashlib
import cv2
# sys.path.append(os.path.join(os.path.dirname(__file__)))
temp_directory = folder_paths.get_temp_directory()
from tqdm import tqdm
import requests
import comfy.utils


CACHE_POOL = {}


class Utils:
    def Md5(str):
        return hashlib.md5(str.encode('utf-8')).hexdigest()

    def check_frames_path(frames_path):

        if frames_path == "" or frames_path.startswith(".") or frames_path.startswith("/") or frames_path.endswith("/") or frames_path.endswith("\\"):
            return "frames_path不能为空"

        frames_path = os.path.join(
            folder_paths.get_output_directory(), frames_path)

        if frames_path == folder_paths.get_output_directory():
            return "frames_path不能为output目录"

        return ""

    def base64_to_pil_image(base64_str):
        if base64_str is None:
            return None
        if len(base64_str) == 0:
            return None
        if type(base64_str) not in [str, bytes]:
            return None
        if base64_str.startswith("data:image/png;base64,"):
            base64_str = base64_str.split(",")[-1]
        base64_str = base64_str.encode("utf-8")
        base64_str = base64.b64decode(base64_str)
        return Image.open(io.BytesIO(base64_str))

    def pil_image_to_base64(pil_image):
        buffered = io.BytesIO()
        pil_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue())
        img_str = str(img_str, encoding="utf-8")
        return f"data:image/png;base64,{img_str}"

    def listdir_png(path):
        try:
            files = os.listdir(path)
            new_files = []
            for file in files:
                if file.endswith(".png"):
                    new_files.append(file)
            files = new_files
            files.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
            return files
        except Exception as e:
            return []

    def listdir_models(path):
        try:
            relative_paths = []
            for root, dirs, files in os.walk(path):
                for file in files:
                    relative_paths.append(os.path.relpath(
                        os.path.join(root, file), path)) 
            relative_paths = [f for f in relative_paths if f.endswith(".safetensors") or f.endswith(
                ".pt") or f.endswith(".pth") or f.endswith(".onnx")]
            return relative_paths

        except Exception as e:

            return []

    def tensor2pil(image):
        return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

    # Convert PIL to Tensor

    def pil2tensor(image):
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)[0]

    def pil2cv(image):
        return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    def cv2pil(image):
        return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    def list_tensor2tensor(data):
        result_tensor = torch.stack(data)
        return result_tensor

    def loadImage(path):
        img = Image.open(path)
        img = img.convert("RGB")
        return img

    def vae_encode_crop_pixels(pixels):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
        return pixels

    def native_vae_encode(vae, image):
        pixels = Utils.vae_encode_crop_pixels(image)
        t = vae.encode(pixels[:, :, :, :3])
        return {"samples": t}

    def native_vae_encode_for_inpaint(vae, pixels, mask):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        mask = torch.nn.functional.interpolate(mask.reshape(
            (-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

        pixels = pixels.clone()
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
            mask = mask[:, :, x_offset:x + x_offset, y_offset:y + y_offset]

        # grow mask by a few pixels to keep things seamless in latent space

        mask_erosion = mask

        m = (1.0 - mask.round()).squeeze(1)
        for i in range(3):
            pixels[:, :, :, i] -= 0.5
            pixels[:, :, :, i] *= m
            pixels[:, :, :, i] += 0.5
        t = vae.encode(pixels)

        return {"samples": t, "noise_mask": (mask_erosion[:, :, :x, :y].round())}

    def native_vae_decode(vae, samples):
        return vae.decode(samples["samples"])

    def native_clip_text_encode(clip, text):
        tokens = clip.tokenize(text)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        return [[cond, {"pooled_output": pooled}]]

    def a1111_clip_text_encode(clip, text):
        try:
            from . import ADV_CLIP_emb_encode
            cond, pooled = ADV_CLIP_emb_encode.advanced_encode(
                clip, text, "none", "A1111", w_max=1.0, apply_to_pooled=False)
            return [[cond, {"pooled_output": pooled}]]
        except Exception as e:
            import nodes
            return nodes.CLIPTextEncode().encode(clip, text)[0]

    def cache_get(key):
        return CACHE_POOL.get(key, None)

    def cache_set(key, value):
        global CACHE_POOL
        CACHE_POOL[key] = value
        return True

    def get_models_path():
        return folder_paths.models_dir

    def get_gguf_models_path():
        models_path = os.path.join(
            folder_paths.models_dir, "gguf")
        os.makedirs(models_path, exist_ok=True)
        return models_path

    def get_translate_object(from_code, to_code):
        try:
            is_disabel_argostranslate = Utils.cache_get(
                "is_disabel_argostranslate")

            if is_disabel_argostranslate is not None:
                return None

            try:
                import argostranslate
                from argostranslate import translate, package
            except ImportError:
                subprocess.run([
                    sys.executable, "-m",
                    "pip", "install", "argostranslate"], check=True)

                try:
                    import argostranslate
                    from argostranslate import translate, package
                except ImportError:
                    Utils.cache_set("is_disabel_argostranslate", True)
                    print(
                        "argostranslate not found and install failed , will disable it")
                    return None

            packages = package.get_installed_packages()
            installed_packages = {}
            for p in packages:
                installed_packages[f"{p.from_code}_{p.to_code}"] = p

            argosmodel_dir = os.path.join(
                Utils.get_models_path(), "argosmodel")
            if not os.path.exists(argosmodel_dir):
                os.makedirs(argosmodel_dir)

            model_name = None
            if from_code == "zh" and to_code == "en":
                model_name = "zh_en"
            elif from_code == "en" and to_code == "zh":
                model_name = "en_zh"
            else:
                return None

            if Utils.cache_get(f"argostranslate_{model_name}") is not None:
                return Utils.cache_get(f"argostranslate_{model_name}")

            if installed_packages.get(model_name, None) is None:
                if not os.path.exists(os.path.join(argosmodel_dir, f"translate-{model_name}-1_9.argosmodel")):
                    argosmodel_file = Utils.download_file(
                        url=f"https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=argosmodel%2Ftranslate-{model_name}-1_9.argosmodel",
                        filepath=os.path.join(
                            argosmodel_dir, f"translate-{model_name}-1_9.argosmodel"),
                    )
                else:
                    argosmodel_file = os.path.join(
                        argosmodel_dir, f"translate-{model_name}-1_9.argosmodel")
                package.install_from_path(argosmodel_file)

            translate_object = translate.get_translation_from_codes(
                from_code=from_code, to_code=to_code)

            Utils.cache_set(f"argostranslate_{model_name}", translate_object)

            return translate_object
        except Exception as e:
            Utils.cache_set("is_disabel_argostranslate", True)
            print(
                "argostranslate not found and install failed , will disable it")
            print(f"get_translate_object error: {e}")
            return None

    def translate_text(text, from_code, to_code):
        translation = Utils.get_translate_object(from_code, to_code)
        if translation is None:
            return text

        # Translate
        translatedText = translation.translate(
            text)

        return translatedText

    def zh2en(text):
        try:
            return Utils.translate_text(text, "zh", "en")
        except Exception as e:
            print(f"zh2en error: {e}")
            return text

    def en2zh(text):
        try:
            return Utils.translate_text(text, "en", "zh")
        except Exception as e:
            print(f"en2zh error: {e}")
            return text

    def prompt_zh_to_en(prompt):
        prompt = prompt.replace("，", ",")
        prompt = prompt.replace("。", ",")
        prompt = prompt.replace("\n", ",")
        tags = prompt.split(",")
        # 判断是否有中文
        for i, tag in enumerate(tags):
            if re.search(u'[\u4e00-\u9fff]', tag):
                tags[i] = Utils.zh2en(tag)
                # 如果第一个字母是大写,转为小写
                if tags[i][0].isupper():
                    tags[i] = tags[i].lower().replace(".", "")

        return ",".join(tags)

    def mask_resize(mask, width, height):
        mask = mask.unsqueeze(0).unsqueeze(0)
        mask = torch.nn.functional.interpolate(
            mask, size=(height, width), mode="bilinear")
        mask = mask.squeeze(0).squeeze(0)
        return mask

    def mask_threshold(interested_mask):
        mask_image = Utils.tensor2pil(interested_mask)
        mask_image_cv2 = Utils.pil2cv(mask_image)
        ret, thresh1 = cv2.threshold(
            mask_image_cv2, 127, 255, cv2.THRESH_BINARY)
        thresh1 = Utils.cv2pil(thresh1)
        thresh1 = np.array(thresh1)
        thresh1 = thresh1[:, :, 0]
        return Utils.pil2tensor(thresh1)

    def mask_erode(interested_mask, value):
        value = int(value)
        mask_image = Utils.tensor2pil(interested_mask)
        mask_image_cv2 = Utils.pil2cv(mask_image)
        kernel = np.ones((5, 5), np.uint8)
        erosion = cv2.erode(mask_image_cv2, kernel, iterations=value)
        erosion = Utils.cv2pil(erosion)
        erosion = np.array(erosion)
        erosion = erosion[:, :, 0]
        return Utils.pil2tensor(erosion)

    def mask_dilate(interested_mask, value):
        value = int(value)
        mask_image = Utils.tensor2pil(interested_mask)
        mask_image_cv2 = Utils.pil2cv(mask_image)
        kernel = np.ones((5, 5), np.uint8)
        dilation = cv2.dilate(mask_image_cv2, kernel, iterations=value)
        dilation = Utils.cv2pil(dilation)
        dilation = np.array(dilation)
        dilation = dilation[:, :, 0]
        return Utils.pil2tensor(dilation)

    def mask_edge_opt(interested_mask, edge_feathering):

        mask_image = Utils.tensor2pil(interested_mask)
        mask_image_cv2 = Utils.pil2cv(mask_image)

        # 高斯模糊
        dilation2 = Utils.cv2pil(mask_image_cv2)
        dilation2 = mask_image.filter(
            ImageFilter.GaussianBlur(edge_feathering))

        # mask_image dilation2 图片蒙版叠加
        dilation2 = Utils.pil2cv(dilation2)
        # dilation2[mask_image_cv2 < 127] = 0
        dilation2 = Utils.cv2pil(dilation2)
        # to RGB
        dilation2 = np.array(dilation2)
        dilation2 = dilation2[:, :, 0]
        return Utils.pil2tensor(dilation2)

    def mask_composite(destination, source, x, y, mask=None, multiplier=8, resize_source=False):
        source = source.to(destination.device)
        if resize_source:
            source = torch.nn.functional.interpolate(source, size=(
                destination.shape[2], destination.shape[3]), mode="bilinear")

        source = comfy.utils.repeat_to_batch_size(source, destination.shape[0])

        x = max(-source.shape[3] * multiplier,
                min(x, destination.shape[3] * multiplier))
        y = max(-source.shape[2] * multiplier,
                min(y, destination.shape[2] * multiplier))

        left, top = (x // multiplier, y // multiplier)
        right, bottom = (left + source.shape[3], top + source.shape[2],)

        if mask is None:
            mask = torch.ones_like(source)
        else:
            mask = mask.to(destination.device, copy=True)
            mask = torch.nn.functional.interpolate(mask.reshape(
                (-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
            mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])

        # calculate the bounds of the source that will be overlapping the destination
        # this prevents the source trying to overwrite latent pixels that are out of bounds
        # of the destination
        visible_width, visible_height = (
            destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)

        mask = mask[:, :, :visible_height, :visible_width]
        inverse_mask = torch.ones_like(mask) - mask

        source_portion = mask * source[:, :, :visible_height, :visible_width]
        destination_portion = inverse_mask * \
            destination[:, :, top:bottom, left:right]

        destination[:, :, top:bottom,
                    left:right] = source_portion + destination_portion
        return destination

    def latent_upscale_by(samples, scale_by):
        s = samples.copy()
        width = round(samples["samples"].shape[3] * scale_by)
        height = round(samples["samples"].shape[2] * scale_by)
        s["samples"] = comfy.utils.common_upscale(
            samples["samples"], width, height, "nearest-exact", "disabled")
        return s

    def resize_by(image, percent):
        # 判断类型是否为PIL
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)

        width, height = image.size
        new_width = int(width * percent)
        new_height = int(height * percent)
        return image.resize((new_width, new_height), Image.LANCZOS)

    def resize_max(im, dst_w, dst_h):
        src_w, src_h = im.size

        if src_h < src_w:
            newWidth = dst_w
            newHeight = dst_w * src_h // src_w
        else:
            newWidth = dst_h * src_w // src_h
            newHeight = dst_h

        newHeight = newHeight // 8 * 8
        newWidth = newWidth // 8 * 8

        return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)

    def get_device():
        return comfy.model_management.get_torch_device()

    def download_small_file(url, filepath):
        response = requests.get(url)
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, "wb") as f:
            f.write(response.content)
        return filepath

    def download_file(url, filepath, threads=8, retries=6):

        get_size_tmp = requests.get(url, stream=True)
        total_size = int(get_size_tmp.headers.get("content-length", 0))

        print(f"Downloading {url} to {filepath} with size {total_size} bytes")

        # 如果文件大小小于 50MB,使用download_small_file
        if total_size < 50 * 1024 * 1024:
            return Utils.download_small_file(url, filepath)

        base_filename = os.path.basename(filepath)
        cache_dir = os.path.join(os.path.dirname(
            filepath), f"{base_filename}.t_{threads}_cache")
        os.makedirs(cache_dir, exist_ok=True)

        def get_total_existing_size():
            fs = os.listdir(cache_dir)
            existing_size = 0
            for f in fs:
                if f.startswith("block_"):
                    existing_size += os.path.getsize(
                        os.path.join(cache_dir, f))
            return existing_size

        total_existing_size = get_total_existing_size()

        if total_size != 0 and total_existing_size != total_size:

            with tqdm(total=total_size, initial=total_existing_size, unit="B", unit_scale=True) as progress_bar:
                all_threads = []

                for i in range(threads):
                    cache_filepath = os.path.join(cache_dir, f"block_{i}")

                    start = total_size // threads * i
                    end = total_size // threads * (i + 1) - 1

                    if i == threads - 1:
                        end = total_size

                    # Check if the file already exists
                    if os.path.exists(cache_filepath):
                        # Get the size of the existing file
                        existing_size = os.path.getsize(cache_filepath)
                    else:
                        existing_size = 0

                    headers = {"Range": f"bytes={start + existing_size}-{end}"}
                    if end == total_size:
                        headers = {"Range": f"bytes={start + existing_size}-"}
                    if start + existing_size >= end:
                        continue
                    # print(f"Downloading {cache_filepath} with headers bytes={start + existing_size}-{end}")

                    # Streaming, so we can iterate over the response.
                    response = requests.get(url, stream=True, headers=headers)

                    def download_file_thread(response, cache_filepath):
                        block_size = 1024
                        if end - (start + existing_size) < block_size:
                            block_size = end - (start + existing_size)
                        with open(cache_filepath, "ab") as file:
                            for data in response.iter_content(block_size):
                                file.write(data)
                                progress_bar.update(
                                    len(data)
                                )

                    t = threading.Thread(
                        target=download_file_thread, args=(response, cache_filepath))

                    all_threads.append(t)

                    t.start()

                for t in all_threads:
                    t.join()

            if total_size != 0 and get_total_existing_size() > total_size:
                # 文件下载失败
                shutil.rmtree(cache_dir)
                raise RuntimeError("Download failed, file is incomplete")

            if total_size != 0 and total_size != get_total_existing_size():
                if retries > 0:
                    retries -= 1
                    print(
                        f"Download failed: {total_size} != {get_total_existing_size()}, retrying... {retries} retries left")
                    return Utils.download_file(url, filepath, threads, retries)

                # 文件损坏
                raise RuntimeError(
                    f"Download failed: {total_size} != {get_total_existing_size()}")

        if os.path.exists(filepath):
            shutil.move(filepath, filepath + ".old." +
                        time.strftime("%Y%m%d%H%M%S"))

        # merge the files
        with open(filepath, "wb") as f:
            for i in range(threads):
                cache_filepath = os.path.join(cache_dir, f"block_{i}")
                with open(cache_filepath, "rb") as cf:
                    f.write(cf.read())

        shutil.rmtree(cache_dir)
        return filepath

    def hf_download_model(url, only_get_path=False):
        if not url.startswith("https://"):
            raise ValueError("URL must start with https://")
        if url.startswith("https://huggingface.co/") or url.startswith("https://hf-mirror.com/"):
            base_model_path = os.path.abspath(os.path.join(
                Utils.get_models_path(), "transformers_models"))
            # https://huggingface.co/FaradayDotDev/llama-3-8b-Instruct-GGUF/resolve/main/llama-3-8b-Instruct.Q2_K.gguf?download=true
            texts = url.split("?")[0].split("/")
            file_name = texts[-1]
            zone_path = f"{texts[3]}/{texts[4]}"

            save_path = os.path.join(base_model_path, zone_path, file_name)

            if os.path.exists(save_path) is False:
                if only_get_path:
                    return None
                os.makedirs(os.path.join(
                    base_model_path, zone_path), exist_ok=True)
                Utils.download_file(url, save_path)

            # Utils.print_log(
            #     f"File {save_path} => {os.path.getsize(save_path)} ")

            # 获取大小
            if os.path.getsize(save_path) == 0:
                if only_get_path:
                    return None
                os.remove(save_path)
                raise ValueError(f"Download failed: {url}")
            return save_path
        else:
            texts = url.split("?")[0].split("/")
            host = texts[2].replace(".", "_")
            base_model_path = os.path.abspath(os.path.join(
                Utils.get_models_path(), f"{host}_models"))

            file_name = texts[-1]
            file_name_no_ext = os.path.splitext(file_name)[0]
            file_ext = os.path.splitext(file_name)[1]
            md5_hash = Utils.Md5(url)

            save_path = os.path.join(
                base_model_path, f"{file_name_no_ext}.{md5_hash}{file_ext}")

            if os.path.exists(save_path) is False:
                if only_get_path:
                    return None
                os.makedirs(base_model_path, exist_ok=True)
                Utils.download_file(url, save_path)

            return save_path

    def print_log(*args):
        if os.environ.get("MZ_DEV", None) is not None:
            print(*args)

    def modelscope_download_model(model_type, model_name, only_get_path=False):
        if model_type not in modelscope_models_map:
            if only_get_path:
                return None
            raise ValueError(f"模型类型 {model_type} 不支持")

        if model_name not in modelscope_models_map[model_type]:
            if only_get_path:
                return None
            error_info = "魔搭可选模型名称列表:\n"
            for key in modelscope_models_map[model_type].keys():
                error_info += f"> {key}\n"
            raise ValueError(error_info)

        model_info = modelscope_models_map[model_type][model_name]
        url = model_info["url"]
        output = model_info["output"]
        save_path = os.path.abspath(
            os.path.join(Utils.get_models_path(), output))
        if not os.path.exists(save_path):
            if only_get_path:
                return None
            save_path = Utils.download_file(url, save_path)
        return save_path

    def progress_bar(steps):
        class pb:
            def __init__(self, steps):
                self.steps = steps
                self.pbar = comfy.utils.ProgressBar(steps)

            def update(self, step, total_steps, pil_img):
                if pil_img is None:
                    self.pbar.update(step, total_steps)

                else:
                    if pil_img.mode != "RGB":
                        pil_img = pil_img.convert("RGB")
                    self.pbar.update_absolute(
                        step, total_steps, ("JPEG", pil_img, 512))

        return pb(steps)

    def split_en_to_zh(text: str):
        if text.find("(") != -1 and text.find(")") != -1:
            sentences = [
                "",
            ]
            for word_index in range(len(text)):
                if text[word_index] == "(" or text[word_index] == ")":
                    sentences.append(str(text[word_index]))
                    sentences.append("")
                else:
                    sentences[-1] += str(text[word_index])

            Utils.print_log("not_translated:", sentences)
            for i in range(len(sentences)):
                if sentences[i] != "(" and sentences[i] != ")":
                    sentences[i] = Utils.split_en_to_zh(sentences[i])

            Utils.print_log("translated:", sentences)

            return "".join(sentences)

        # 中文标点转英文标点
        text = text.replace("，", ",")
        text = text.replace("。", ".")
        text = text.replace("？", "?")
        text = text.replace("！", "!")
        text = text.replace("；", ";")

        result = []
        if text.find("\n") != -1:
            text = text.split("\n")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return "\n".join(result)

        if text.find(".") != -1:
            text = text.split(".")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return ".".join(result)

        if text.find("?") != -1:
            text = text.split("?")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return "?".join(result)

        if text.find("!") != -1:
            text = text.split("!")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return "!".join(result)

        if text.find(";") != -1:
            text = text.split(";")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return ";".join(result)

        if text.find(",") != -1:
            text = text.split(",")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return ",".join(result)

        if text.find(":") != -1:
            text = text.split(":")
            for t in text:
                if t != "":
                    result.append(Utils.split_en_to_zh(t))
                else:
                    result.append(t)
            return ":".join(result)

        # 如果是纯数字,不翻译
        if text.isdigit() or text.replace(".", "").isdigit() or text.replace(" ", "").isdigit() or text.replace("-", "").isdigit():
            return text

        return Utils.en2zh(text)

    def to_debug_prompt(p):
        if p is None:
            return ""
        zh = Utils.en2zh(p)
        if p == zh:
            return p
        zh = Utils.split_en_to_zh(p)
        p = p.strip()
        return f"""
原文:
{p}

中文翻译:
{zh}
"""

    def get_gguf_files():
        gguf_dir = Utils.get_gguf_models_path()
        if not os.path.exists(gguf_dir):
            os.makedirs(gguf_dir)
        gguf_files = []
        # walk gguf_dir
        for root, dirs, files in os.walk(gguf_dir):
            for file in files:
                if file.endswith(".gguf"):
                    gguf_files.append(
                        os.path.relpath(os.path.join(root, file), gguf_dir))

        return gguf_files

    def get_comfyui_models_path():
        return folder_paths.models_dir

    def download_model(model_info, only_get_path=False):

        url = model_info["url"]
        output = model_info["output"]
        save_path = os.path.abspath(
            os.path.join(Utils.get_comfyui_models_path(), output))
        if not os.path.exists(save_path):
            if only_get_path:
                return None
            save_path = Utils.download_file(url, save_path)
        return save_path

    def file_hash(file_path, hash_method):
        if not os.path.isfile(file_path):
            return ''
        h = hash_method()
        with open(file_path, 'rb') as f:
            while b := f.read(8192):
                h.update(b)
        return h.hexdigest()

    def get_cache_by_local(key):
        try:
            cache_json_file = os.path.join(
                Utils.get_models_path(), f"caches.json")

            if not os.path.exists(cache_json_file):
                return None

            with open(cache_json_file, "r", encoding="utf-8") as f:
                cache_json = json.load(f)
                return cache_json.get(key, None)
        except:
            return None

    def set_cache_by_local(key, value):
        try:
            cache_json_file = os.path.join(
                Utils.get_models_path(), f"caches.json")

            if not os.path.exists(cache_json_file):
                cache_json = {}
            else:
                with open(cache_json_file, "r", encoding="utf-8") as f:
                    cache_json = json.load(f)

            cache_json[key] = value

            with open(cache_json_file, "w", encoding="utf-8") as f:
                json.dump(cache_json, f, indent=4)
        except:
            pass

    def file_sha256(file_path):
        # 获取文件的更新时间
        file_stat = os.stat(file_path)
        file_mtime = file_stat.st_mtime
        file_size = file_stat.st_size
        cache_key = f"{file_path}_{file_mtime}_{file_size}"
        cache_value = Utils.get_cache_by_local(cache_key)
        if cache_value is not None:
            return cache_value

        sha256 = Utils.file_hash(file_path, hashlib.sha256)
        Utils.set_cache_by_local(cache_key, sha256)
        return sha256

    def get_auto_model_fullpath(model_name):
        fullpath = Utils.cache_get(f"get_auto_model_fullpath_{model_name}")
        Utils.print_log(f"get_auto_model_fullpath_{model_name} => {fullpath}")
        if fullpath is not None:
            if os.path.exists(fullpath):
                return fullpath

        find_paths = []
        target_sha256 = ""
        file_path = ""
        download_url = ""

        MODEL_ZOO = Utils.get_model_zoo()
        for model in MODEL_ZOO:
            if model["model"] == model_name:
                find_paths = model["find_path"]
                target_sha256 = model["SHA256"]
                file_path = model["file_path"]
                download_url = model["url"]
                break

        if target_sha256 == "":
            raise ValueError(f"Model {model_name} not found in MODEL_ZOO")

        if os.path.exists(file_path):
            if Utils.file_sha256(file_path) != target_sha256:
                print(f"Model {model_name} file hash not match...")
            return file_path

        for find_path in find_paths:
            find_fullpath = os.path.join(
                Utils.get_comfyui_models_path(), find_path)

            if os.path.exists(find_fullpath):
                for root, dirs, files in os.walk(find_fullpath):
                    for file in files:
                        if target_sha256 == Utils.file_sha256(os.path.join(root, file)):
                            Utils.cache_set(
                                f"get_auto_model_fullpath_{model_name}", os.path.join(root, file))
                            return os.path.join(root, file)
                        else:
                            Utils.print_log(
                                f"Model {os.path.join(root, file)} file hash not match, {target_sha256} != {Utils.file_sha256(os.path.join(root, file))}")

        result = Utils.download_model(
            {"url": download_url, "output": file_path})
        Utils.cache_set(f"get_auto_model_fullpath_{model_name}", result)
        return result

    def testDownloadSpeed(url):
        try:
            print(f"Testing download speed for {url}")
            start = time.time()
            # 下载2M数据
            headers = {"Range": "bytes=0-2097151"}
            _ = requests.get(url, headers=headers, timeout=5)
            end = time.time()
            print(
                f"Download speed: {round(5.00 / (float(end) - float(start)) / 1024, 2)} KB/s")
            return float(end) - float(start) < 4
        except Exception as e:
            print(f"Test download speed failed: {e}")
            return False

    def get_model_zoo(tags_filter=None):
        source_model_zoo_file = os.path.join(
            os.path.dirname(__file__), "configs", "model_zoo.json")
        source_model_zoo_json = []
        try:
            with open(source_model_zoo_file, "r", encoding="utf-8") as f:
                source_model_zoo_json = json.load(f)
        except:
            pass

        # Utils.print_log(f"source_model_zoo_json: {json.dumps(source_model_zoo_json, indent=4)}")
        if tags_filter is not None:
            source_model_zoo_json = [
                m for m in source_model_zoo_json if tags_filter in m["tags"]]

        return source_model_zoo_json
        


modelscope_models_map = {
   
}