arc2face_nodes.py

import os
import logging
import torch
import numpy as np
import math
from PIL import Image
import cv2
import folder_paths
from insightface.app import FaceAnalysis
from insightface.utils import storage
from diffusers import StableDiffusionPipeline, UNet2DConditionModel, DPMSolverMultistepScheduler, AutoPipelineForImage2Image
import json
from safetensors.torch import load_file
from .utils import project_face_embs
from .models import CLIPTextModelWrapper

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Arc2FaceFaceExtractor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "average_method": (["average", "median", "trimmed_mean", "ensemble_average", "ensemble_median", "max_pooling", "min_pooling", "rounded_mode", "rounded_mode_averaging", "random_sampling"],),
                "n_outliers": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1}),
            }
        }

    RETURN_TYPES = ("FACE_EMBEDDING",)
    FUNCTION = "extract_face_embedding"
    CATEGORY = "Arc2Face"

    def __init__(self):
        def custom_download(sub_dir, name, force, root='~/.insightface'):
            return os.path.join(folder_paths.models_dir, "antelopev2")

        storage.download = custom_download
        self.app = FaceAnalysis(name='antelopev2', root=os.path.join(folder_paths.models_dir, ""), providers=['CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))

    def extract_face_embedding(self, images, average_method, n_outliers):
        try:
            img_np = images.cpu().numpy()
            if img_np.ndim == 3:
                img_np = np.expand_dims(img_np, axis=0)
            if img_np.ndim != 4:
                raise ValueError(f"Invalid image dimensions: {img_np.shape}")
            
            all_embeddings = []
            
            for img in img_np:
                img = np.clip(255. * img, 0, 255).astype(np.uint8)
                if img.shape[2] == 1:
                    img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
                elif img.shape[2] == 4:
                    img = img[:, :, :3]
                elif img.shape[2] != 3:
                    raise ValueError(f"Unexpected number of channels: {img.shape[2]}")

                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

                min_size = 64
                if img.shape[0] < min_size or img.shape[1] < min_size:
                    scale = min_size / min(img.shape[0], img.shape[1])
                    img = cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)

                faces = self.app.get(img)
                
                if faces:
                    embeddings = [torch.tensor(face.embedding, dtype=torch.float32) for face in faces]
                    all_embeddings.extend(embeddings)
            
            if not all_embeddings:
                raise ValueError("No faces detected in any of the images")

            all_embeddings = self.remove_outliers(all_embeddings, n_outliers)
            avg_embedding = self.average_embeddings(all_embeddings, method=average_method)
            
            return (avg_embedding,)
        
        except Exception as e:
            logger.error(f"Error in Arc2FaceFaceExtractor: {str(e)}")
            return (torch.zeros(512, dtype=torch.float32),)

    def remove_outliers(self, embeddings, n_outliers):
        if n_outliers == 0 or n_outliers >= len(embeddings):
            return embeddings

        centroid = torch.mean(torch.stack(embeddings), dim=0)
        distances = [torch.norm(e - centroid, p=2).item() for e in embeddings]
        outlier_indices = sorted(range(len(distances)), key=lambda i: distances[i], reverse=True)[:n_outliers]
        return [e for i, e in enumerate(embeddings) if i not in outlier_indices]

    def average_embeddings(self, embeddings, method="average"):
        if not embeddings:
            raise ValueError("No valid embeddings to average")

        logger.info(f"Combining {len(embeddings)} embeddings using method: {method}")

        embeddings_stack = torch.stack(embeddings)

        if method == "average":
            return torch.mean(embeddings_stack, dim=0)
        elif method == "median":
            return torch.median(embeddings_stack, dim=0).values
        elif method == "trimmed_mean":
            lower_bound, upper_bound = int(0.15 * len(embeddings)), int(0.85 * len(embeddings))
            sorted_embeddings = embeddings_stack.sort(dim=0).values
            trimmed_embeddings = sorted_embeddings[lower_bound:upper_bound, :]
            return torch.mean(trimmed_embeddings, dim=0)
        elif method == "max_pooling":
            return torch.max(embeddings_stack, dim=0).values
        elif method == "min_pooling":
            return torch.min(embeddings_stack, dim=0).values
        elif method == "rounded_mode":
            mode_embeddings = torch.zeros_like(embeddings_stack[0])
            for dim in range(embeddings_stack.size(1)):
                current_dim_values = embeddings_stack[:, dim]
                rounded_dim_values = torch.round(current_dim_values * 100) / 100
                values, counts = rounded_dim_values.unique(return_counts=True)
                mode_embeddings[dim] = values[counts.argmax()]
            return mode_embeddings
        elif method == "rounded_mode_averaging":
            averaged_mode_embeddings = torch.zeros_like(embeddings_stack[0])
            for dim in range(embeddings_stack.size(1)):
                current_dim_values = embeddings_stack[:, dim]
                rounded_dim_values = torch.round(current_dim_values * 100) / 100
                mode_value, _ = rounded_dim_values.mode()
                contributing_indices = (rounded_dim_values == mode_value).nonzero(as_tuple=True)[0]
                contributing_values = current_dim_values[contributing_indices]
                averaged_mode_embeddings[dim] = torch.mean(contributing_values)
            return averaged_mode_embeddings
        elif method == "ensemble_average":
            mean_embedding = torch.mean(embeddings_stack, dim=0)
            median_embedding = torch.median(embeddings_stack, dim=0).values
            lower_bound, upper_bound = int(0.15 * len(embeddings)), int(0.85 * len(embeddings))
            sorted_embeddings = embeddings_stack.sort(dim=0).values
            trimmed_embeddings = sorted_embeddings[lower_bound:upper_bound, :]
            trimmed_mean_embedding = torch.mean(trimmed_embeddings, dim=0)
            all_averages = torch.stack([mean_embedding, median_embedding, trimmed_mean_embedding])
            return torch.mean(all_averages, dim=0)
        elif method == "ensemble_median":
            mean_embedding = torch.mean(embeddings_stack, dim=0)
            median_embedding = torch.median(embeddings_stack, dim=0).values
            lower_bound, upper_bound = int(0.15 * len(embeddings)), int(0.85 * len(embeddings))
            sorted_embeddings = embeddings_stack.sort(dim=0).values
            trimmed_embeddings = sorted_embeddings[lower_bound:upper_bound, :]
            trimmed_mean_embedding = torch.mean(trimmed_embeddings, dim=0)
            all_averages = torch.stack([mean_embedding, median_embedding, trimmed_mean_embedding])
            return torch.median(all_averages, dim=0).values
        elif method == "random_sampling":
            randomly_sampled_embedding = torch.empty(embeddings_stack.shape[1], dtype=embeddings_stack.dtype).to(embeddings_stack.device)
            for dim in range(embeddings_stack.shape[1]):
                random_index = torch.randint(0, embeddings_stack.shape[0], (1,)).item()
                randomly_sampled_embedding[dim] = embeddings_stack[random_index, dim]
            return randomly_sampled_embedding
        else:
            raise ValueError("Unsupported averaging method.")

class Arc2FaceUNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"model_path": ("STRING", {"default": "diffusion_pytorch_model.safetensors"})}}

    RETURN_TYPES = ("ARC2FACE_UNET",)
    FUNCTION = "load_unet"
    CATEGORY = "Arc2Face"

    def load_unet(self, model_path):
        arc2face_path = os.path.join(folder_paths.models_dir, "arc2face_checkpoints")
        full_path = os.path.join(arc2face_path, model_path)
        
        with open(os.path.join(arc2face_path, "config.json"), 'r') as f:
            config = json.load(f)

        unet = UNet2DConditionModel(**config)
        unet_state_dict = load_file(full_path)
        unet.load_state_dict(unet_state_dict)
        return (unet,)

class Arc2FaceEncoderLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"encoder_path": ("STRING", {"default": "encoder"})}}

    RETURN_TYPES = ("ARC2FACE_ENCODER",)
    FUNCTION = "load_encoder"
    CATEGORY = "Arc2Face"

    def load_encoder(self, encoder_path):
        arc2face_path = os.path.join(folder_paths.models_dir, "arc2face_checkpoints")
        full_path = os.path.join(arc2face_path, encoder_path)
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        encoder = CLIPTextModelWrapper.from_pretrained(
            full_path,
            torch_dtype=dtype
        )
        return (encoder,)

class Arc2FaceGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "face_embedding": ("FACE_EMBEDDING",),
            "unet": ("ARC2FACE_UNET",),
            "encoder": ("ARC2FACE_ENCODER",),
            "negative_prompt": ("STRING", {"default": "ugly, deformed, noisy, blurry, low contrast, split image"}),
            "num_inference_steps": ("INT", {"default": 30, "min": 1, "max": 100}),
            "guidance_scale": ("FLOAT", {"default": 2.7, "min": 0.1, "max": 30.0, "step": 0.1}),
            "num_images": ("INT", {"default": 1, "min": 1, "max": 4}),
            "width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
            "height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
            "seed": ("INT", {"default": -1}),
        }}

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "Arc2Face"

    def __init__(self):
        self.pipe = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32

    def load_pipeline(self, unet, encoder):
        if self.pipe is None:
            base_model = 'botp/stable-diffusion-v1-5'
            
            self.pipe = StableDiffusionPipeline.from_pretrained(
                base_model,
                text_encoder=encoder,
                unet=unet,
                torch_dtype=self.dtype,
                safety_checker=None
            )
            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
            self.pipe = self.pipe.to(self.device)

    def generate(self, face_embedding, unet, encoder, negative_prompt, num_inference_steps, guidance_scale, num_images, width, height, seed):
        self.load_pipeline(unet, encoder)

        if seed == -1:
            seed = torch.randint(0, 2**32 - 1, (1,)).item()
        generator = torch.Generator(self.device).manual_seed(seed)

        face_embedding = face_embedding.to(dtype=torch.float16, device=self.device)
        face_embedding = face_embedding / torch.norm(face_embedding, dim=0, keepdim=True)

        with torch.autocast(device_type=self.device, dtype=torch.float16):
            id_emb = project_face_embs(self.pipe, face_embedding.unsqueeze(0))
            
            output = self.pipe(
                prompt_embeds=id_emb,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                num_images_per_prompt=num_images,
                width=width,
                height=height,
                generator=generator,
            )

        images = output.images
        comfy_images = []
        for img in images:
            img = img.convert("RGB")
            img = np.array(img).astype(np.float32) / 255.0
            img = torch.from_numpy(img)[None,]
            comfy_images.append(img)

        if len(comfy_images) == 1:
            return (comfy_images[0],)
        return (torch.cat(comfy_images, dim=0),)

class Arc2FaceImg2ImgGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "face_embedding": ("FACE_EMBEDDING",),
            "unet": ("ARC2FACE_UNET",),
            "encoder": ("ARC2FACE_ENCODER",),
            "initial_image": ("IMAGE",),
            "negative_prompt": ("STRING", {"default": "ugly, deformed, noisy, blurry, low contrast, split image"}),
            "num_inference_steps": ("INT", {"default": 30, "min": 1, "max": 100}),
            "guidance_scale": ("FLOAT", {"default": 2.7, "min": 0.1, "max": 30.0, "step": 0.1}),
            "num_images": ("INT", {"default": 1, "min": 1, "max": 4}),
            "seed": ("INT", {"default": -1}),
            "denoise_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
            "extra_param": ("STRING", {"default": ""})
        }}

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "Arc2Face"

    def __init__(self):
        self.pipe = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32

    def load_pipeline(self, unet, encoder):
        if self.pipe is None:
            base_model = 'botp/stable-diffusion-v1-5'
            
            self.pipe = AutoPipelineForImage2Image.from_pretrained(
                base_model,
                text_encoder=encoder,
                unet=unet,
                torch_dtype=self.dtype,
                safety_checker=None
            )
            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
            self.pipe = self.pipe.to(self.device)

    def generate(self, face_embedding, unet, encoder, initial_image, negative_prompt, num_inference_steps, guidance_scale, num_images, seed, denoise_strength, extra_param):
        self.load_pipeline(unet, encoder)

        logger.info(f"Initial image shape: {initial_image.shape}")

        if seed == -1:
            seed = torch.randint(0, 2**32 - 1, (1,)).item()
        generator = torch.Generator(self.device).manual_seed(seed)

        face_embedding = face_embedding.to(dtype=torch.float16, device=self.device)
        face_embedding = face_embedding / torch.norm(face_embedding, dim=0, keepdim=True)

        initial_image_np = initial_image.squeeze().cpu().numpy()
        logger.info(f"Initial image shape after squeeze: {initial_image_np.shape}")
        
        initial_image_pil = Image.fromarray((initial_image_np * 255).astype(np.uint8), mode='RGB')

        logger.info(f"Initial image size: {initial_image_pil.size}")

        width, height = initial_image_pil.size

        if width < 64 or height < 64:
            raise ValueError(f"Initial image is too small, width is {width} and height is {height}")
        
        if initial_image_pil is None:
            raise ValueError("No initial image provided")

        with torch.autocast(device_type=self.device, dtype=torch.float16):
            id_emb = project_face_embs(self.pipe, face_embedding.unsqueeze(0), additional_text=extra_param)
            
            output = self.pipe(
                image=initial_image_pil,
                prompt_embeds=id_emb,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                num_images_per_prompt=num_images,
                generator=generator,
                strength=denoise_strength,
            )

        images = output.images
        comfy_images = []
        for img in images:
            img = img.convert("RGB")
            img = np.array(img).astype(np.float32) / 255.0
            img = torch.from_numpy(img)[None,]
            comfy_images.append(img)

        if len(comfy_images) == 1:
            return (comfy_images[0],)
        return (torch.cat(comfy_images, dim=0),)

class Arc2FaceImageGridGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "directory": ("STRING", {"default": ""}),
                "max_images": ("INT", {"default": 16, "min": 1, "max": 64, "step": 1}),
                "max_size": ("INT", {"default": 768, "min": 64, "max": 2048, "step": 64}),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate_grid"
    CATEGORY = "Arc2Face"

    def resize_image(self, img, max_size):
        width, height = img.size
        if width > height:
            if width > max_size:
                height = int(max_size * height / width)
                width = max_size
        else:
            if height > max_size:
                width = int(max_size * width / height)
                height = max_size
        return img.resize((width, height), Image.LANCZOS)

    def generate_grid(self, directory, max_images, max_size):
        image_files = [f for f in os.listdir(directory) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
        
        if not image_files:
            raise ValueError("No image files found in the specified directory.")
        
        image_files = image_files[:max_images]
        
        images = [self.resize_image(Image.open(os.path.join(directory, img)).convert('RGB'), max_size) for img in image_files]
        
        n = len(images)
        cols = math.ceil(math.sqrt(n))
        rows = math.ceil(n / cols)
        
        grid_width = cols * max_size
        grid_height = rows * max_size
        grid_img = Image.new('RGB', (grid_width, grid_height))
        
        for i, img in enumerate(images):
            x = (i % cols) * max_size
            y = (i // cols) * max_size
            grid_img.paste(img, (x, y))
        
        grid_np = np.array(grid_img).astype(np.float32) / 255.0
        grid_tensor = torch.from_numpy(grid_np)[None,]
        
        logger.info(f"Grid tensor shape: {grid_tensor.shape}")
        logger.info(f"Grid tensor dtype: {grid_tensor.dtype}")
        
        return (grid_tensor,)

NODE_CLASS_MAPPINGS = {
    "Arc2FaceFaceExtractor": Arc2FaceFaceExtractor,
    "Arc2FaceUNetLoader": Arc2FaceUNetLoader,
    "Arc2FaceEncoderLoader": Arc2FaceEncoderLoader,
    "Arc2FaceGenerator": Arc2FaceGenerator,
    "Arc2FaceImg2ImgGenerator": Arc2FaceImg2ImgGenerator,
    "Arc2FaceImageGridGenerator": Arc2FaceImageGridGenerator
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Arc2FaceFaceExtractor": "Arc2Face Face Extractor",
    "Arc2FaceUNetLoader": "Arc2Face UNet Loader",
    "Arc2FaceEncoderLoader": "Arc2Face Encoder Loader",
    "Arc2FaceGenerator": "Arc2Face Generator",
    "Arc2FaceImg2ImgGenerator": "Arc2Face Img2Img Generator",
    "Arc2FaceImageGridGenerator": "Arc2Face Image Grid Generator"
}