From 1ed629a5675aa412c11d230444e3ffded8495e9f Mon Sep 17 00:00:00 2001 From: zhangsongbo Date: Mon, 9 Sep 2024 20:25:42 +0800 Subject: [PATCH] Fixed a bug and added support for more face images --- comfy_extras/nodes_photomaker_v2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_photomaker_v2.py b/comfy_extras/nodes_photomaker_v2.py index b314add98..f1fa11761 100644 --- a/comfy_extras/nodes_photomaker_v2.py +++ b/comfy_extras/nodes_photomaker_v2.py @@ -236,7 +236,7 @@ class PhotoMakerEncodeV2: if index > 0: id_embeds = self.get_id_embeds(image) token_index = index - 1 - num_id_images = 1 + num_id_images = image.shape[0] num_tokens_v2 = 2 id_embeds = id_embeds.unsqueeze(0).to(device=photomaker.load_device) class_tokens_mask = [True if token_index <= i < token_index+(num_id_images*num_tokens_v2) else False for i in range(cond.shape[1])] @@ -248,7 +248,7 @@ class PhotoMakerEncodeV2: return ([[out, {"pooled_output": pooled}]], [[cond, {"pooled_output": pooled}]],) - def tensor2np(tensor: torch.Tensor): + def tensor2np(self, tensor: torch.Tensor): if len(tensor.shape) == 3: # Single image return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8) else: # Batch of images