# !/usr/bin/env python # -*- coding: UTF-8 -*- import os import random import sys import yaml import torch import cv2 from PIL import Image import numpy as np from transformers import CLIPTextModel, CLIPTokenizer from omegaconf import OmegaConf from safetensors import safe_open from insightface.app import FaceAnalysis from insightface.utils import face_align from diffusers import (AutoencoderKL, DDIMScheduler, ControlNetModel, KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler, AutoPipelineForInpainting, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler, UNet2DConditionModel, KDPM2DiscreteScheduler, EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, StableDiffusionXLControlNetPipeline, DDPMScheduler, TCDScheduler, LCMScheduler) from .faceadapter.face_adapter import FaceAdapterPlusForVideoLora from .animatediff.utils.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, \ convert_ldm_vae_checkpoint from .animatediff.utils.util import load_weights from .animatediff.pipelines.pipeline_animation import AnimationPipeline from .animatediff.models.unet import UNet3DConditionModel from .animatediff.utils.util import save_videos_grid, export_to_video import folder_paths dir_path = os.path.dirname(os.path.abspath(__file__)) path_dir = os.path.dirname(dir_path) file_path = os.path.dirname(path_dir) yamlPath = os.path.join(file_path, "extra_model_paths.yaml") if os.path.isfile(yamlPath): f = open(yamlPath, 'r', encoding='utf-8') d1 = yaml.load(f,Loader=yaml.SafeLoader) a111 = d1['a111'] if d1['a111'] is not None else [] if a111 != "": other_path =a111['base_path'] if a111['base_path'] is not None else '' other_checkpoint = a111['checkpoints'] if a111['checkpoints'] is not None else '' other_model_path = os.path.join(other_path, other_checkpoint) else: other_model_path = os.path.join(file_path,"models","checkpoints") paths = [] for search_path in folder_paths.get_folder_paths("diffusers"): if os.path.exists(search_path): for root, subdir, files in os.walk(search_path, followlinks=True): if "model_index.json" in files: paths.append(os.path.relpath(root, start=search_path)) if paths != []: paths = [] + [x for x in paths if x] else: paths = ["no model in default diffusers directory", ] scheduler_list = [ "DDIM", "DDPM", "DPM++ 2M", "DPM++ 2M Karras", "DPM++ 2M SDE", "DPM++ 2M SDE Karras", "DPM++ SDE", "DPM++ SDE Karras", "DPM2", "DPM2 Karras", "DPM2 a", "DPM2 a Karras", "Heun", "LCM", "LMS", "LMS Karras", "UniPC", "UniPC_Bh2", "TCD" ] def get_sheduler(name): scheduler = False if name == "DDIM": scheduler = DDIMScheduler elif name == "DDPM": scheduler = DDPMScheduler elif name == "DPM++ 2M": scheduler = DPMSolverMultistepScheduler elif name == "DPM++ 2M Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True) elif name == "DPM++ 2M SDE": scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++") elif name == "DPM++ 2M SDE Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++") elif name == "DPM++ SDE": scheduler = DPMSolverSinglestepScheduler elif name == "DPM++ SDE Karras": scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True) elif name == "DPM2": scheduler = KDPM2DiscreteScheduler elif name == "DPM2 Karras": scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True) elif name == "DPM2 a": scheduler = KDPM2AncestralDiscreteScheduler elif name == "DPM2 a Karras": scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True) elif name == "Heun": scheduler = HeunDiscreteScheduler elif name == "LCD": scheduler = LCMScheduler elif name == "LMS": scheduler = LMSDiscreteScheduler elif name == "LMS Karras": scheduler = LMSDiscreteScheduler(use_karras_sigmas=True) elif name == "UniPC_Bh1": scheduler = UniPCMultistepScheduler(solver_type="bh1") elif name == "UniPC_Bh2": scheduler = UniPCMultistepScheduler(solver_type="bh2") elif name == "TCD": scheduler = TCDScheduler return scheduler def get_local_path(file_path, model_path): path = os.path.join(file_path, "models", "diffusers", model_path) model_path = os.path.normpath(path) if sys.platform.startswith('win32'): model_path = model_path.replace('\\', "/") return model_path def get_instance_path(path): os_path = os.path.normpath(path) if sys.platform.startswith('win32'): os_path = os_path.replace('\\', "/") return os_path class ID_Animator: def __init__(self): pass def tensor_to_image(self, tensor): tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "prompt": ("STRING", {"multiline": True, "default": "Iron Man soars through the clouds, his repulsors blazing"}), "negative_prompt": ("STRING", {"multiline": True, "default": "semi-realistic, cgi, 3d, render, sketch, cartoon," " drawing, anime, text, close up, cropped, out of frame," " worst quality, low quality, jpeg artifacts, ugly, duplicate," " morbid, mutilated, extra fingers, mutated hands, poorly drawn hands," " poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, " "bad proportions, extra limbs, cloned face, disfigured, gross proportions," " malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers," " too many fingers, long neck"}), "model_local_path": (paths,), "repo_id": ("STRING", {"default": "runwayml/stable-diffusion-v1-5"}), "scheduler": (scheduler_list,), "checkpoints_list": (folder_paths.get_filename_list("checkpoints"),), "steps": ("INT", {"default": 30, "min": 1, "max": 2048, "step": 1, "display": "number"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "cfg": ("FLOAT", {"default": 8, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64, "display": "number"}), "width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64, "display": "number"}), "video_length": ("INT", {"default": 16, "min": 1, "max": 100}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("text",) FUNCTION = "id_animator" CATEGORY = "ID_Animator" def load_model(self, inference_config, sd_version, scheduler, id_ckpt, image_encoder_path, dreambooth_model_path, motion_module_path): inference_config = OmegaConf.load(inference_config) tokenizer = CLIPTokenizer.from_pretrained(sd_version, subfolder="tokenizer", torch_dtype=torch.float16, ) text_encoder = CLIPTextModel.from_pretrained(sd_version, subfolder="text_encoder", torch_dtype=torch.float16, ).cuda() vae = AutoencoderKL.from_pretrained(sd_version, subfolder="vae", torch_dtype=torch.float16, ).cuda() unet = UNet3DConditionModel.from_pretrained_2d(sd_version, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container( inference_config.unet_additional_kwargs) ).cuda() scheduler_used = get_sheduler(scheduler) pipeline = AnimationPipeline( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, controlnet=None, # beta_start=0.00085, beta_end=0.012, beta_schedule="linear",steps_offset=1 scheduler=scheduler_used(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs) # scheduler=EulerAncestralDiscreteScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs) # scheduler=EulerAncestralDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="linear",steps_offset=1 ), torch_dtype=torch.float16, ).to("cuda") pipeline = load_weights( pipeline, # motion module motion_module_path=motion_module_path, motion_module_lora_configs=[], # domain adapter adapter_lora_path="", adapter_lora_scale=1, # image layers dreambooth_model_path=None, lora_model_path="", lora_alpha=0.8 ).to("cuda") if dreambooth_model_path != "": print(f"load dreambooth model from {dreambooth_model_path}") dreambooth_state_dict = {} with safe_open(dreambooth_model_path, framework="pt", device="cpu") as f: for key in f.keys(): dreambooth_state_dict[key] = f.get_tensor(key) converted_vae_checkpoint = convert_ldm_vae_checkpoint(dreambooth_state_dict, pipeline.vae.config) # print(vae) # vae ->to_q,to_k,to_v # print(converted_vae_checkpoint) convert_vae_keys = list(converted_vae_checkpoint.keys()) for key in convert_vae_keys: if "encoder.mid_block.attentions" in key or "decoder.mid_block.attentions" in key: new_key = None if "key" in key: new_key = key.replace("key", "to_k") elif "query" in key: new_key = key.replace("query", "to_q") elif "value" in key: new_key = key.replace("value", "to_v") elif "proj_attn" in key: new_key = key.replace("proj_attn", "to_out.0") if new_key: converted_vae_checkpoint[new_key] = converted_vae_checkpoint.pop(key) pipeline.vae.load_state_dict(converted_vae_checkpoint) converted_unet_checkpoint = convert_ldm_unet_checkpoint(dreambooth_state_dict, pipeline.unet.config) pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False) pipeline.text_encoder = convert_ldm_clip_checkpoint(dreambooth_state_dict).to("cuda") del dreambooth_state_dict pipeline = pipeline.to(torch.float16) id_animator = FaceAdapterPlusForVideoLora(pipeline, image_encoder_path, id_ckpt, num_tokens=16, device=torch.device("cuda"), torch_type=torch.float16) return id_animator def id_animator(self, image, prompt, negative_prompt, model_local_path, repo_id, scheduler, checkpoints_list,steps,seed, cfg, height, width, video_length): if model_local_path == ["no model in default diffusers directory", ] and repo_id == "": raise "you need fill repo_id or download model in diffusers dir " model_path = get_local_path(file_path, model_local_path) if repo_id == "": repo_id = model_path inference_config = f"{dir_path}/inference-v2.yaml" sd_version = repo_id path_ckpt = os.path.join(dir_path, "models", "animator.ckpt") path_dream = os.path.join(other_model_path,checkpoints_list) path_motion = os.path.join(dir_path, "models/animatediff_models/mm_sd_v15_v2.ckpt") path_img = os.path.join(dir_path, "models/image_encoder") # print(inference_config, path_dream, path_ckpt, path_motion) id_ckpt = get_instance_path(path_ckpt) image_encoder_path = get_instance_path(path_img) dreambooth_model_path = get_instance_path(path_dream) motion_module_path = get_instance_path(path_motion) text_encode_global = globals() app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) app.prepare(ctx_id=0, det_size=(320, 320)) animator = self.load_model(inference_config, sd_version, scheduler, id_ckpt, image_encoder_path, dreambooth_model_path, motion_module_path) Pil_img = self.tensor_to_image(image) img = cv2.cvtColor(np.asarray(Pil_img), cv2.COLOR_RGB2BGR) faces = app.get(img) face_roi = face_align.norm_crop(img, faces[0]['kps'], 112) face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB) pil_image = [Image.fromarray(face_roi).resize((224, 224))] sample = animator.generate(pil_image, negative_prompt=negative_prompt, prompt=prompt, num_inference_steps=steps, seed=seed, guidance_scale=cfg, width=width, height=height, video_length=video_length, scale=0.8, ) filename_prefix = ''.join(random.choice("0123456789") for _ in range(6))+".gif" output = os.path.join(file_path, "output", filename_prefix) save_videos_grid(sample, output) # export_to_video(sample, f"{output}") return (output,) NODE_CLASS_MAPPINGS = { "ID_Animator": ID_Animator, } NODE_DISPLAY_NAME_MAPPINGS = { "ID_Animator": "ID_Animator" }