import json import os import shutil import subprocess import sys import threading import time import traceback from typing import Tuple import warnings 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 = Utils.listdir(path) new_files = [] for file in files: if file.lower().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(path): try: files = os.listdir(path) # 排除.开头的文件 files = [file for file in files if not file.startswith(".")] return files except Exception as e: return [] def tensor2pil(image): return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) def tensors2pil_list(images): return [Utils.tensor2pil(image) for image in images] # 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 cache_get(key): return CACHE_POOL.get(key, None) def cache_set(key, value): global CACHE_POOL CACHE_POOL[key] = value return True def model_cache_clean(model_type): global CACHE_POOL for key in list(CACHE_POOL.keys()): if key.startswith(f"model_cache_{model_type}"): del CACHE_POOL[key] torch.cuda.empty_cache() return True def model_cache_get(model_type, model_path): resp = Utils.cache_get(f"model_cache_{model_type}") if resp is None: return None cache_model_path = resp.get("model_path") if cache_model_path != model_path: Utils.model_cache_clean(f"model_cache_{model_type}") return None return resp.get("model") def model_cache_set(model_type, model_path, model): key = f"model_cache_{model_type}" if Utils.cache_get(key) is not None: Utils.model_cache_clean(key) Utils.cache_set(key, {"model_path": model_path, "model": model}) return True def get_minus_zone_models_path(): models_path = os.path.join( folder_paths.models_dir, "minus_zone_models") os.makedirs(models_path, exist_ok=True) return models_path def get_comfyui_models_path(): return folder_paths.models_dir def translate_text(text, from_code, to_code): try: import argostranslate from argostranslate import translate except ImportError: subprocess.run([ sys.executable, "-m", "pip", "install", "argostranslate"], check=True) import argostranslate from argostranslate import translate try: translation = translate.get_translation_from_codes( from_code, to_code) if translation is None: raise Exception("Translation not found") except Exception as e: print(e) argostranslate.package.update_package_index() available_packages = argostranslate.package.get_available_packages() package_to_install = next( filter( lambda x: (x.from_code == from_code and x.to_code == to_code), available_packages, ) ) download_path = package_to_install.download() print("package_to_install.download():", download_path) argostranslate.package.install_from_path(download_path) translation = translate.get_translation_from_codes( from_code, to_code) if translation is None: return text # Translate translatedText = translation.translate( text) return translatedText def zh2en(text): return Utils.translate_text(text, "zh", "en") def en2zh(text): return Utils.translate_text(text, "en", "zh") 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 resize_min(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 add_watermark(image, watermark): if watermark == "": return image try: import PIL from PIL import ImageDraw, ImageFont except ImportError: subprocess.run([ sys.executable, "-m", "pip", "install", "Pillow"], check=True) import PIL from PIL import ImageDraw, ImageFont # 获取PIL版本号 pil_version = PIL.__version__ if pil_version >= "10.0.0": def textsize(self, text, font): left, top, right, bottom = self.textbbox((0, 0), text, font) return right - left, bottom - top ImageDraw.ImageDraw.textsize = textsize font_fullpath = Utils.download_model( { "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=font%2FAlibabaPuHuiTi-2-75-SemiBold.ttf", "output": "font/AlibabaPuHuiTi-2-75-SemiBold.ttf", } ) watermarks = watermark.split("\n") width, height = image.size short_edge = min(width, height) font_size = short_edge // 12 font = ImageFont.truetype(font_fullpath, font_size) # print("pil_version:", pil_version) draw = ImageDraw.Draw(image) text = watermarks[0] textwidth, textheight = draw.textsize(text, font) x = (width - textwidth) // 2 bottom = 10 y = height - textheight - (textheight * 0.4 + bottom + 8) draw.text((x, y), text, font=font) if len(watermarks) > 1: y1 = y + textheight text = watermarks[1] font_size = int(font_size * 0.4) font = ImageFont.truetype(font_fullpath, font_size) textwidth, textheight = draw.textsize(text, font) x = (width - textwidth) // 2 y = y1 - bottom + 4 draw.text((x, y), text, font=font) return image def get_device(): return comfy.model_management.get_torch_device() 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") 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 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 load_lora(model, lora_path, strength_model): lora = comfy.utils.load_torch_file(lora_path, safe_load=True) model_lora, _ = comfy.sd.load_lora_for_models( model, None, lora, strength_model, 0) return model_lora def load_checkpoint(ckpt_name): cache_data = Utils.cache_get(ckpt_name) if cache_data is not None: print("load from cache: ", ckpt_name) return cache_data ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, output_clip=True) model, clip, vae = out[:3] Utils.cache_set(ckpt_name, (model, clip, vae)) return model, clip, vae 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 file_sha256(file_path): return Utils.file_hash(file_path, hashlib.sha256) def get_auto_model_fullpath(model_name): find_paths = [] target_sha256 = "" file_path = "" download_url = "" 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( folder_paths.get_output_directory(), find_path) if os.path.exists(find_fullpath): for root, dirs, files in os.walk(find_fullpath): # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for file in files: if target_sha256 == Utils.file_sha256(os.path.join(root, file)): return os.path.join(root, file) return Utils.download_model({"url": download_url, "output": file_path}) def is_sd15_model(model): type_str = str(type(model.model.model_config).__name__) return "SD15" in type_str def is_sdxl_model(model): type_str = str(type(model.model.model_config).__name__) return "SDXL" in type_str def progress_bar(steps, taesd_type="sd1_5"): class pb: if taesd_type == "sd1_5": taesd_decoder_name = "taesd_decoder" elif taesd_type == "sdxl": taesd_decoder_name = "taesdxl_decoder" latent_rgb_factors = None latent_channels = 4 def __init__(self, steps): self.steps = steps self.pbar = comfy.utils.ProgressBar(steps) def get_previewer(self): import latent_preview previewer = latent_preview.get_previewer( 'cuda', self) return previewer def update(self, step, total_steps, pil_img=None): try: pil_img_info = ("JPEG", pil_img, 512) if pil_img is None: pil_img_info = None if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: pil_img_info = pil_img # print("pil_img_info:", type(pil_img), pil_img_info) # print("step:", step, "total_steps:", total_steps) self.pbar.update_absolute( step, total_steps, pil_img_info) except Exception as e: print("progress_bar:", e) raise e return pb(steps) def get_free_port(): import socket s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(('localhost', 0)) port = s.getsockname()[1] s.close() return port def Simple_Server(reader): import threading from http.server import HTTPServer, BaseHTTPRequestHandler port = Utils.get_free_port() is_running = {"value": True} def stop_server(): is_running.update({"value": False}) httpd = None class SimpleHTTPRequestHandler(BaseHTTPRequestHandler): def do_POST(self): self.send_response(200) self.end_headers() content_length = int(self.headers['Content-Length']) post_data = self.rfile.read(content_length) data = json.loads(post_data) if reader(data) == False: stop_server() def log_message(self, format, *args): pass httpd = HTTPServer(('localhost', port), SimpleHTTPRequestHandler) def serve_forever(): try: print( "===========================httpd.serve_forever() start=======================================") httpd.serve_forever() except Exception as e: pass print( "===========================httpd.serve_forever() end=======================================") threading.Thread(target=serve_forever).start() def check_server_stop(): while is_running["value"]: # print("is_running : ", is_running) time.sleep(1) httpd.shutdown() httpd.server_close() threading.Thread(target=check_server_stop).start() return stop_server, port def xy_image(pre_render_images, pre_render_texts_x, pre_render_texts_y): # 去掉pre_render_texts_x中所有相同字符串前缀 def get_common_prefix(pre_render_texts_x): if len(pre_render_texts_x) == 0: return "" common_prefix = "" for i in range(len(pre_render_texts_x[0])): c = pre_render_texts_x[0][i] for j in range(len(pre_render_texts_x)): if pre_render_texts_x[j][i] != c: return common_prefix else: common_prefix += c return common_prefix common_prefix = get_common_prefix(pre_render_texts_x) if common_prefix != "": pre_render_texts_x = [ x.replace(common_prefix, "") for x in pre_render_texts_x] x_enable_num = len(pre_render_images) y_enable_num = len(pre_render_images[0]) max_width = 0 max_height = 0 for x in range(0, len(pre_render_images)): for y in range(0, len(pre_render_images[x])): if pre_render_images[x][y].width < 512: org_width = pre_render_images[x][y].width org_height = pre_render_images[x][y].height pre_render_images[x][y] = pre_render_images[x][y].resize( (512, int(org_height * 512 / org_width))) if pre_render_images[x][y].width > max_width: max_width = pre_render_images[x][y].width if pre_render_images[x][y].height > max_height: max_height = pre_render_images[x][y].height image_xy_canvas = Image.new( "RGB", (max_width * x_enable_num, max_height * y_enable_num)) for i in range(x_enable_num): for j in range(y_enable_num): image_xy_canvas.paste( pre_render_images[i][j], (max_width * i, max_height * j)) # draw axis from PIL import ImageDraw, ImageFont import PIL # 获取PIL版本号 pil_version = PIL.__version__ if pil_version >= "10.0.0": def textsize(self, text, font): left, top, right, bottom = self.textbbox((0, 0), text, font) return right - left, bottom - top ImageDraw.ImageDraw.textsize = textsize full_padding = 260 full_canvas_width = max_width * x_enable_num + full_padding * 2 full_canvas_height = max_height * y_enable_num + full_padding * 2 full_canvas = Image.new( "RGB", (full_canvas_width, full_canvas_height), "white") full_canvas.paste(image_xy_canvas, (full_padding, full_padding)) draw = ImageDraw.Draw(full_canvas) # top draw.line((full_padding, full_padding, full_canvas_width - full_padding, full_padding), fill="black", width=4) # left draw.line((full_padding, full_padding, full_padding, full_canvas_height - full_padding), fill="black", width=4) # bottom draw.line((full_padding, full_canvas_height - full_padding, full_canvas_width - full_padding, full_canvas_height - full_padding), fill="black", width=4) # right draw.line((full_canvas_width - full_padding, full_padding, full_canvas_width - full_padding, full_canvas_height - full_padding), fill="black", width=4) font_fullpath = Utils.download_model( { "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=font%2FAlibabaPuHuiTi-2-75-SemiBold.ttf", "output": "font/AlibabaPuHuiTi-2-75-SemiBold.ttf", } ) if os.path.exists(font_fullpath): font = ImageFont.truetype(font_fullpath, size=32,) else: font = ImageFont.load_default() for i in range(x_enable_num): textwidth, textheight = draw.textsize( pre_render_texts_x[i], font) offset_x = max_width * i + full_padding + \ ((max_width - textwidth) // 2) offset_y = full_padding - textheight - 24 draw.text((offset_x, offset_y), pre_render_texts_x[i], font=font, fill="black") for j in range(y_enable_num): textwidth, textheight = draw.textsize( pre_render_texts_y[j], font) label_text = pre_render_texts_y[j] def is_exceed_width(t): textwidth, _ = draw.textsize( t, font) return textwidth > max_width - 24 if textwidth > full_padding - 24: # 超过宽度就换行 label_text = "" for c in pre_render_texts_y[j]: if is_exceed_width(label_text + c): break label_text += c offset_x = full_padding - textwidth - 24 offset_y = max_height * j + full_padding + \ ((max_height - textheight) // 2) draw.text((offset_x, offset_y), label_text, font=font, fill="black") return full_canvas def get_models_by_folder(dir_path): models = [] for root, dirs, files in os.walk(dir_path): # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for file in files: if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"): models.append(os.path.join(root, file)) return models def get_folders_by_folder(dir_path): folders = [] for root, dirs, files in os.walk(dir_path): # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for dir in dirs: folders.append(os.path.join(root, dir)) return folders from torch import nn class CustomizeEmbedsModel(nn.Module): dtype = torch.float16 shared = None # x = torch.zeros(1, 1, 256, 2048) x = None def __init__(self, *args, **kwargs): super().__init__() def to(self, *args, **kwargs): return self def forward(self, *args, **kwargs): # print("CustomizeEmbedsModel forward: args:", args) # print("CustomizeEmbedsModel forward: kwargs:", kwargs) input_ids = kwargs.get("input_ids", None) # if self.x is None: if True: if input_ids is None: batch_size = 1 else: batch_size = input_ids.shape[0] attention_mask = kwargs.get("attention_mask") attention_mask_dim = attention_mask.shape[1] self.x = torch.zeros(1, batch_size, 256, 2048, dtype=self.dtype) if kwargs.get("output_hidden_states", False): return { "hidden_states": self.x.to("cuda"), "input_ids": torch.zeros(1, 1), } return self.x class CustomizeTokenizer(dict): added_tokens_encoder = [] input_ids = None attention_mask = None batch_size = 1 def __init__(self, *args, **kwargs): self['added_tokens_encoder'] = self.added_tokens_encoder self['input_ids'] = self.input_ids self['attention_mask'] = self.attention_mask self.batch_size = kwargs.get("batch_size", 1) def tokenize(self, text): return text def __call__(self, *args, **kwargs): # print("CustomizeTokenizer args:", args) # print("CustomizeTokenizer kwargs:", kwargs) value = args[0] if isinstance(value, str): batch_size = 1 else: batch_size = value.shape[0] # print(f"CustomizeTokenizer batch_size: {batch_size}") # if self.input_ids is not None: # return self self.input_ids = torch.zeros(batch_size, 256) self.attention_mask = torch.zeros(batch_size, 256) self['input_ids'] = self.input_ids self['attention_mask'] = self.attention_mask # print("CustomizeTokenizer input_ids:", self.input_ids.shape) # print("CustomizeTokenizer attention_mask:", self.attention_mask.shape) return self class CustomizeEmbeds(): def __init__(self): super().__init__() self.tokenizer = CustomizeTokenizer() self.model = CustomizeEmbedsModel().to("cuda") self.max_length = 256 class CustomizeMT5Embedder(nn.Module): device = torch.device("cuda") def __init__( self, model_dir="t5-v1_1-xxl", model_kwargs=None, torch_dtype=None, use_tokenizer_only=False, max_length=128, batch_size=1, ): super().__init__() self.torch_dtype = torch_dtype or torch.bfloat16 self.max_length = max_length self.tokenizer = CustomizeTokenizer( batch_size=batch_size ) self.model = CustomizeEmbedsModel().to("cuda") def gradient_checkpointing_enable(self): pass def gradient_checkpointing_disable(self): pass def get_tokens_and_mask(self, texts): text_tokens_and_mask = self.tokenizer( texts, max_length=self.max_length, padding="max_length", truncation=True, return_attention_mask=True, add_special_tokens=True, return_tensors="pt", ) tokens = text_tokens_and_mask["input_ids"][0] mask = text_tokens_and_mask["attention_mask"][0] return tokens, mask def get_text_embeddings(self, texts, attention_mask=True, layer_index=-1): text_tokens_and_mask = self.tokenizer( texts, max_length=self.max_length, padding="max_length", truncation=True, return_attention_mask=True, add_special_tokens=True, return_tensors="pt", ) outputs = self.model( input_ids=text_tokens_and_mask["input_ids"], attention_mask=( text_tokens_and_mask["attention_mask"] if attention_mask else None ), output_hidden_states=True, ) text_encoder_embs = outputs["hidden_states"][layer_index].detach() return text_encoder_embs, text_tokens_and_mask["attention_mask"].to(self.device) def get_input_ids(self, caption): return self.tokenizer( caption, padding="max_length", truncation=True, max_length=self.max_length, return_tensors="pt", ).input_ids def get_hidden_states(self, input_ids, layer_index=-1): return self.get_text_embeddings(input_ids, layer_index=layer_index) class HSubprocess: process_instance = None process_instance_pid = None screen_name = None _mswindows = False def __init__(self, args, screen_name=None): self._mswindows = (sys.platform == "win32") if self._mswindows: self.args = ["start", "/w"] self.args.extend(args) return self.screen_name = screen_name if screen_name is not None: try: subprocess.check_call(["screen", "-v"]) except Exception as e: raise Exception("Please install screen first.") screen_cmd = ["screen", "-R", screen_name, "-m", ] screen_cmd.extend(args) self.args = screen_cmd else: self.args = args def stop(self): if self._mswindows: print('taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') os.system( f'taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') if self.process_instance is not None: if self.screen_name is not None: subprocess.run( ["screen", "-d", self.screen_name]) self.process_instance.kill() self.process_instance = None try: try: import psutil except ImportError: subprocess.check_call( [sys.executable, "-m", "pip", "install", "psutil"]) import psutil psutil.Process(self.process_instance_pid).terminate() except Exception as e: print(e) self.process_instance_pid = None def wait(self): with subprocess.Popen( self.args, stdin=subprocess.PIPE, shell=True, ) as process: self.process_instance = process self.process_instance_pid = process.pid print(f"Subprocess PID: {self.process_instance_pid}") try: stdout, stderr = process.communicate() except subprocess.TimeoutExpired as exc: process.kill() if self._mswindows: exc.stdout, exc.stderr = process.communicate() else: process.wait() raise except Exception as e: process.wait() raise retcode = process.poll() if retcode != 0: raise subprocess.CalledProcessError(retcode, process.args) self.process_instance = None self.process_instance_pid = None