# !/usr/bin/env python # -*- coding: UTF-8 -*- import os import torch import gc from PIL import Image import numpy as np import math import comfy.utils import cv2 import random import torchaudio import folder_paths from comfy.utils import common_upscale,ProgressBar from safetensors.torch import load_file import soundfile as sf import comfy.model_management as mm from pathlib import PureWindowsPath cur_path = os.path.dirname(os.path.abspath(__file__)) def covert_obj_img(images, masks,target_size): background_color=(255, 255, 255) if images is not None and masks is not None: output_list = [] # Convert to numpy for easier processing images_np = images.cpu().numpy() # Handle mask shapes if masks.dim() == 2: # Single mask HW, expand to BHW masks_np = masks.unsqueeze(0).cpu().numpy() elif masks.dim() == 3: # Batch masks BHW masks_np = masks.cpu().numpy() else: raise ValueError("Mask must be of shape HW or BHW") assert masks_np.shape[0] == images_np.shape[0] , "Masks and images must had same batch size" batch_size = images_np.shape[0] for i in range(batch_size): # Get image and corresponding mask img = images_np[i] # HWC mask = masks_np[i] # Ensure image is in range [0, 255] if img.max() <= 1.0: img = img * 255 img = img.astype(np.uint8) # Create RGBA image using mask as alpha channel rgba_img = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8) rgba_img[:, :, :3] = img rgba_img[:, :, 3] = ((1 - mask) * 255).astype(np.uint8) # # Convert to PIL Image pil_img = Image.fromarray(rgba_img, 'RGBA') pil_img.save(f"{i}temp.png") output_list.append(pil_img) img=output_list[0] original_width, original_height = img.size target_width, target_height = target_size # 1. 计算缩放比例，确保图片能完整放入目标框内 ratio = min(target_width / original_width, target_height / original_height) # 2. 计算缩放后的新尺寸 new_width = int(original_width * ratio) new_height = int(original_height * ratio) # 3. 高质量缩放图片 resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) # 4. 创建一个新的纯色背景画布 # 注意：颜色需要是RGBA格式，所以白色是 (255, 255, 255, 255) # 最后一个值255代表完全不透明 final_background_color = background_color + (255,) background = Image.new('RGBA', target_size, final_background_color) # 5. 计算粘贴位置，使其居中 paste_x = (target_width - new_width) // 2 paste_y = (target_height - new_height) // 2 # 6. 将缩放后的图片粘贴到背景画布上 # 第三个参数 `resized_img` 作为蒙版，可以正确处理PNG的透明通道 background.paste(resized_img, (paste_x, paste_y), resized_img) else: final_background_color = background_color + (255,) background = Image.new('RGBA', target_size, final_background_color) background.save("temp.png") return background.convert('RGB') def get_wav2vec_repo(repo): required_files = ["chinese-wav2vec2-base-fairseq-ckpt.pt", "config.json", "model.safetensors", "preprocessor_config.json"] if not repo: wav2vec_repo=os.path.join(folder_paths.models_dir, "wav2vec2-base") if not os.path.exists(wav2vec_repo): os.makedirs(wav2vec_repo) download_file(folder_paths.models_dir, required_files) else: if not check_files_exist(wav2vec_repo,required_files): download_file(folder_paths.models_dir, required_files) else: wav2vec_repo=PureWindowsPath(repo).as_posix() return wav2vec_repo def download_file(local_dir, required_files): from huggingface_hub import hf_hub_download for i in required_files: hf_hub_download( repo_id="youliang1233214/InteractAvatar", subfolder="wav2vec2-base", filename=i, local_dir = local_dir, ) def check_files_exist(folder_path, required_files): for file_name in required_files: file_path = os.path.join(folder_path, file_name) if not os.path.exists(file_path): return False return True def clear_comfyui_cache(): cf_models=mm.loaded_models() try: for pipe in cf_models: pipe.unpatch_model(device_to=torch.device("cpu")) print(f"Unpatching models.{pipe}") except: pass mm.soft_empty_cache() torch.cuda.empty_cache() max_gpu_memory = torch.cuda.max_memory_allocated() print(f"After Max GPU memory allocated: {max_gpu_memory / 1000 ** 3:.2f} GB") # def trans2path(audio): # if audio is None: # return None # import io as io_base # audio_file_prefix = ''.join(random.choice("0123456789") for _ in range(6)) # audio_file = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}_temp.wav") # buff = io_base.BytesIO() # torchaudio.save(buff, audio["waveform"].squeeze(0), audio["sample_rate"], format="FLAC") # with open(audio_file, 'wb') as f: # f.write(buff.getbuffer()) # return audio_file def trans2path(audio): """ 修正版：使用 soundfile 代替 torchaudio.save 以避开 torchcodec 的环境报错。 """ if audio is None: return None import io as io_base audio_file_prefix = ''.join(random.choice("0123456789") for _ in range(6)) audio_file = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}_temp.wav") buff = io_base.BytesIO() # --- 修正逻辑开始 --- # ComfyUI 音频格式通常为 [Batch, Channels, Samples] -> [1, C, S] # 我们需要将其转换为 NumPy，并调整维度为 soundfile 要求的 [Samples, Channels] waveform = audio["waveform"].squeeze(0).cpu().numpy() # 结果为 [C, S] sample_rate = audio["sample_rate"] if waveform.ndim == 2: waveform = waveform.T # 转置为 [S, C] # 使用 soundfile 直接写入内存流，不触发 torchaudio 的后端检测 sf.write(buff, waveform, sample_rate, format="FLAC") # --- 修正逻辑结束 --- with open(audio_file, 'wb') as f: f.write(buff.getbuffer()) return audio_file def encode_image( image, vae): if image is None: return None ref_latents=None samples = image.movedim(-1, 1) total = int(1024 * 1024) scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) width = round(samples.shape[3] * scale_by) height = round(samples.shape[2] * scale_by) s = comfy.utils.common_upscale(samples, width, height, "area", "disabled") image = s.movedim(1, -1) if vae is not None: ref_latents = vae.encode(image[:, :, :, :3]) return ref_latents def add_mean(latents): vae_config={"latents_mean": [ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 ], "latents_std": [ 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916 ],} latents_mean = (torch.tensor(vae_config["latents_mean"]).view(1, 16, 1, 1, 1).to(latents.device, latents.dtype)) latents_std = 1.0 / torch.tensor(vae_config["latents_std"]).view(1, 16, 1, 1, 1).to(latents.device, latents.dtype) latents = latents / latents_std + latents_mean image_latent_height, image_latent_width = latents.shape[3:] image_latents = pack_latents_( latents, 1, 16, image_latent_height, image_latent_width) return image_latents def pack_latents_(latents, batch_size, num_channels_latents, height, width): latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) latents = latents.permute(0, 2, 4, 1, 3, 5) latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) return latents def load_lora(model, lora_1, lora_2, lora_scale1, lora_scale2): lora_path_1=folder_paths.get_full_path("loras", lora_1) if lora_1 != "none" else None lora_path_2=folder_paths.get_full_path("loras", lora_2) if lora_2 != "none" else None # lora_list=[i for i in [lora_path_1,lora_path_2] if i is not None] # lora_scales=[lora_scale1,lora_scale2] all_adapters = model.get_list_adapters() dit_list=[] if all_adapters: dit_list= all_adapters.get('transformer',[])+all_adapters.get('transformer_2',[]) if lora_path_1 is not None: adapter_name=os.path.splitext(os.path.basename(lora_path_1))[0].replace(".", "_") dit_list2=all_adapters.get('transformer_2',[]) if dit_list2: if adapter_name in dit_list: #dit_list pass else: for i in dit_list2: model.delete_adapters(i) print(f"去除dit中未加载的lora: {i}") try: model.load_lora_weights(lora_path_1, adapter_name=adapter_name,**{"load_into_transformer_2": True}) model.set_adapters([adapter_name], adapter_weights=lora_scale1) except KeyError as e: try: print(f"检测到特殊的 LoRA 格式，尝试手动处理: {lora_path_1}") state_dict = torch.load(lora_path_1, map_location="cpu",weights_only=False) if not lora_path_1.endswith(".safetensors") else load_file(lora_path_1,) processed_state_dict = preprocess_lora_state_dict(state_dict) model.load_lora_weights(processed_state_dict, adapter_name=adapter_name,**{"load_into_transformer_2": True}) model.set_adapters([adapter_name], adapter_weights=lora_scale1) except: print(f"加载LoRA权重失败: {e}") pass else: try: model.load_lora_weights(lora_path_1, adapter_name=adapter_name,**{"load_into_transformer_2": True}) model.set_adapters([adapter_name], adapter_weights=lora_scale1) except KeyError as e: try: print(f"检测到特殊的 LoRA 格式，尝试手动处理: {lora_path_1}") state_dict = torch.load(lora_path_1, map_location="cpu",weights_only=False) if not lora_path_1.endswith(".safetensors") else load_file(lora_path_1,) processed_state_dict = preprocess_lora_state_dict(state_dict) model.load_lora_weights(processed_state_dict, adapter_name=adapter_name,**{"load_into_transformer_2": True}) model.set_adapters([adapter_name], adapter_weights=lora_scale1) del processed_state_dict except: print(f"加载LoRA权重失败: {e}") pass if lora_path_2 is not None: adapter_name=os.path.splitext(os.path.basename(lora_path_2))[0].replace(".", "_") dit_list=all_adapters.get('transformer',[]) if dit_list: if adapter_name in dit_list: #dit_list pass else: for i in dit_list: model.delete_adapters(i) print(f"去除dit中未加载的lora: {i}") try: model.load_lora_weights(lora_path_2, adapter_name=adapter_name,**{"load_into_transformer_2": False}) model.set_adapters([adapter_name], adapter_weights=lora_scale2) except KeyError as e: try: print(f"检测到特殊的 LoRA 格式，尝试手动处理: {lora_path_2}") state_dict = torch.load(lora_path_2, map_location="cpu",weights_only=False) if not lora_path_2.endswith(".safetensors") else load_file(lora_path_2,) processed_state_dict = preprocess_lora_state_dict(state_dict) model.load_lora_weights(processed_state_dict, adapter_name=adapter_name,**{"load_into_transformer_2": False}) model.set_adapters([adapter_name], adapter_weights=lora_scale2) del processed_state_dict except: print(f"加载LoRA权重失败: {e}") pass else: try: model.load_lora_weights(lora_path_2, adapter_name=adapter_name,**{"load_into_transformer_2": False}) model.set_adapters([adapter_name], adapter_weights=lora_scale2) except KeyError as e: try: print(f"检测到特殊的 LoRA 格式，尝试手动处理: {lora_path_2}") state_dict = torch.load(lora_path_2, map_location="cpu",weights_only=False) if not lora_path_2.endswith(".safetensors") else load_file(lora_path_2,) processed_state_dict = preprocess_lora_state_dict(state_dict) model.load_lora_weights(processed_state_dict, adapter_name=adapter_name,**{"load_into_transformer_2": False}) model.set_adapters([adapter_name], adapter_weights=lora_scale2) del processed_state_dict except: print(f"加载LoRA权重失败: {e}") pass return model def preprocess_lora_state_dict(state_dict): processed_dict = state_dict.copy() keys_to_remove = [ 'head.head.diff_b', 'head.head.diff_m', 'head.head.diff', 'patch_embedding.diff', 'patch_embedding.diff_b', 'blocks.*.diff_m', # 匹配所有blocks的diff_m 'head.head.lora_down' 'diffusion_model.head.head.diff' 'diffusion_model.head.head.diff_b' 'diffusion_model.head.lora_down' ] keys_to_delete = [] for key in processed_dict.keys(): if key.endswith('.diff_m'): keys_to_delete.append(key) for key in keys_to_delete: processed_dict.pop(key, None) print(f"移除键: {key}") for key in keys_to_remove: if key in processed_dict: processed_dict.pop(key, None) print(f"移除键: {key}") return processed_dict def gc_cleanup(): gc.collect() torch.cuda.empty_cache() def tensor2cv(tensor_image): if len(tensor_image.shape)==4:# b hwc to hwc tensor_image=tensor_image.squeeze(0) if tensor_image.is_cuda: tensor_image = tensor_image.cpu() tensor_image=tensor_image.numpy() #反归一化 maxValue=tensor_image.max() tensor_image=tensor_image*255/maxValue img_cv2=np.uint8(tensor_image)#32 to uint8 img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR) return img_cv2 def phi2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img def tensor2image(tensor): tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def tensor2pillist(tensor_in): d1, _, _, _ = tensor_in.size() if d1 == 1: img_list = [tensor2image(tensor_in)] else: tensor_list = torch.chunk(tensor_in, chunks=d1) img_list=[tensor2image(i) for i in tensor_list] return img_list def tensor2pillist_upscale(tensor_in,width,height): d1, _, _, _ = tensor_in.size() if d1 == 1: img_list = [nomarl_upscale(tensor_in,width,height)] else: tensor_list = torch.chunk(tensor_in, chunks=d1) img_list=[nomarl_upscale(i,width,height) for i in tensor_list] return img_list def tensor2list(tensor_in,width,height): if tensor_in is None: return None d1, _, _, _ = tensor_in.size() if d1 == 1: tensor_list = [tensor_upscale(tensor_in,width,height)] else: tensor_list_ = torch.chunk(tensor_in, chunks=d1) tensor_list=[tensor_upscale(i,width,height) for i in tensor_list_] return tensor_list def tensor_upscale(tensor, width, height): samples = tensor.movedim(-1, 1) samples = common_upscale(samples, width, height, "bilinear", "center") samples = samples.movedim(1, -1) return samples def nomarl_upscale(img, width, height): samples = img.movedim(-1, 1) img = common_upscale(samples, width, height, "bilinear", "center") samples = img.movedim(1, -1) img = tensor2image(samples) return img def cv2tensor(img,bgr2rgb=True): assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) if bgr2rgb: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img.transpose((2, 0, 1))) return img.float().div(255).permute(1, 2, 0).unsqueeze(0) def images_generator(img_list: list, ): # get img size sizes = {} for image_ in img_list: if isinstance(image_, Image.Image): count = sizes.get(image_.size, 0) sizes[image_.size] = count + 1 elif isinstance(image_, np.ndarray): count = sizes.get(image_.shape[:2][::-1], 0) sizes[image_.shape[:2][::-1]] = count + 1 else: raise "unsupport image list,must be pil or cv2!!!" size = max(sizes.items(), key=lambda x: x[1])[0] yield size[0], size[1] # any to tensor def load_image(img_in): if isinstance(img_in, Image.Image): img_in = img_in.convert("RGB") i = np.array(img_in, dtype=np.float32) i = torch.from_numpy(i).div_(255) if i.shape[0] != size[1] or i.shape[1] != size[0]: i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0) i = common_upscale(i, size[0], size[1], "lanczos", "center") i = i.squeeze(0).movedim(0, -1).numpy() return i elif isinstance(img_in, np.ndarray): i = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB).astype(np.float32) i = torch.from_numpy(i).div_(255) print(i.shape) return i else: raise "unsupport image list,must be pil,cv2 or tensor!!!" total_images = len(img_list) processed_images = 0 pbar = ProgressBar(total_images) images = map(load_image, img_list) try: prev_image = next(images) while True: next_image = next(images) yield prev_image processed_images += 1 pbar.update_absolute(processed_images, total_images) prev_image = next_image except StopIteration: pass if prev_image is not None: yield prev_image def load_images_list(img_list: list, ): gen = images_generator(img_list) (width, height) = next(gen) images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3))))) if len(images) == 0: raise FileNotFoundError(f"No images could be loaded .") return images def get_video_files(directory, extensions=None): if extensions is None: extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov'] extensions = [ext.lower() for ext in extensions] video_files = [] for root, dirs, files in os.walk(directory): for file in files: _, ext = os.path.splitext(file) ext = ext.lower()[1:] if ext in extensions: full_path = os.path.join(root, file) video_files.append(full_path) return video_files