fixed the Plus model

This commit is contained in:
matt3o
2023-09-07 18:18:02 +02:00
parent aec92d8faa
commit 034fb6b4f2
3 changed files with 693 additions and 64 deletions
+77 -64
View File
@@ -3,10 +3,10 @@ import os
import comfy.utils
import comfy.model_management
from comfy.clip_vision import ClipVisionModel
from torch import nn
import torch.nn.functional as F
from transformers import CLIPVisionModelWithProjection
from .resampler import Resampler
@@ -75,12 +75,10 @@ def attention(q, k, v, extra_options):
return out
class IPAdapter(nn.Module):
def __init__(self, ipadapter_model, clip_embeddings, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
def __init__(self, ipadapter_model, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
super().__init__()
self.ipadapter_model = ipadapter_model
self.clip_embeddings = clip_embeddings
self.clip_embeddings_dim = clip_embeddings_dim
self.cross_attention_dim = self.ipadapter_model["ip_adapter"]["1.to_k_ip.weight"].shape[1]
self.clip_extra_context_tokens = clip_extra_context_tokens
@@ -100,9 +98,9 @@ class IPAdapter(nn.Module):
return image_proj_model
@torch.inference_mode()
def get_image_embeds(self, clip_image_embeds):
image_prompt_embeds = self.image_proj_model(clip_image_embeds)
uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(clip_image_embeds))
def get_image_embeds(self, clip_embed, clip_embed_zeroed):
image_prompt_embeds = self.image_proj_model(clip_embed)
uncond_image_prompt_embeds = self.image_proj_model(clip_embed_zeroed)
return image_prompt_embeds, uncond_image_prompt_embeds
class IPAdapterPlus(IPAdapter):
@@ -120,10 +118,9 @@ class IPAdapterPlus(IPAdapter):
return image_proj_model
@torch.inference_mode()
def get_image_embeds(self, clip_image_embeds):
image_prompt_embeds = self.image_proj_model(clip_image_embeds)
uncond_clip_image_embeds = self.clip_vision(torch.zeros_like(self.clip_embeddings), output_hidden_states=True).hidden_states[-2]
uncond_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds)
def get_image_embeds(self, clip_embed, clip_embed_zeroed):
image_prompt_embeds = self.image_proj_model(clip_embed)
uncond_image_prompt_embeds = self.image_proj_model(clip_embed_zeroed)
return image_prompt_embeds, uncond_image_prompt_embeds
class CrossAttentionPatch:
@@ -175,6 +172,21 @@ class CrossAttentionPatch:
return out.to(dtype=org_dtype)
class IPAdapterCLIPVisionEncode:
@classmethod
def INPUT_TYPES(s):
return {"required": { "clip_vision": ("CLIP_VISION",),
"image": ("IMAGE",)
}}
RETURN_TYPES = ("CLIP_VISION_OUTPUT",)
FUNCTION = "encode"
CATEGORY = "ipadapter"
def encode(self, clip_vision, image):
output = clip_vision.encode_image_IPAdapter(image)
return (output,)
class IPAdapterModelLoader:
@classmethod
def INPUT_TYPES(s):
@@ -183,7 +195,7 @@ class IPAdapterModelLoader:
RETURN_TYPES = ("IPADAPTER",)
FUNCTION = "load_ipadapter_model"
CATEGORY = "loaders"
CATEGORY = "ipadapter"
def load_ipadapter_model(self, ipadapter_file):
ckpt_path = os.path.join(MODELS_DIR, ipadapter_file)
@@ -213,49 +225,58 @@ class IPAdapterApply:
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_ipadapter"
CATEGORY = "conditioning/ipadapter"
CATEGORY = "ipadapter"
def apply_ipadapter(self, model, ipadapter, clip_vision_output, weight, mask=None):
self.dtype = model.model.diffusion_model.dtype
self.device = comfy.model_management.get_torch_device()
self.weight = weight
clip_embeddings = clip_vision_output.image_embeds.to(self.device, dtype=self.dtype)
clip_embeddings_dim = clip_embeddings.shape[-1]
# check if we are using the 'CLIP Vision Encode IPAdapter' node
self.is_ipadapter_encode = type(clip_vision_output) is tuple
if self.is_ipadapter_encode:
clip_embed, clip_embed_zeroed = clip_vision_output
else:
clip_embed = clip_vision_output.image_embeds
clip_embed_zeroed = torch.zeros_like(clip_embed)
self.is_plus = "latents" in ipadapter["image_proj"]
cross_attention_dim = ipadapter["ip_adapter"]["1.to_k_ip.weight"].shape[1]
if self.is_plus:
clip_extra_context_tokens = ipadapter["image_proj"]["latents"].shape[1]
# IPAdapter Plus requires the IPAdapter Encoder
if not self.is_ipadapter_encode:
raise Exception("IPAdapter Plus requires 'CLIP Vision Encode (IPAdapter)' node")
clip_extra_context_tokens = 16
clip_embed = clip_embed.hidden_states[-2]
else:
clip_extra_context_tokens = 4 # ipadapter["image_proj"]["proj.weight"].shape[0] <-- doesn't work?
# if we come 'CLIP Vision Encode IPAdapter' we need to extract the image_embeds
if self.is_ipadapter_encode:
clip_embed = clip_embed.image_embeds
clip_embed_zeroed = torch.zeros_like(clip_embed)
clip_extra_context_tokens = 4
clip_embeddings_dim = clip_embed.shape[-1]
self.is_sdxl = cross_attention_dim == 2048
if self.is_plus:
self.ipadapter = IPAdapterPlus(
ipadapter,
clip_embeddings,
cross_attention_dim=cross_attention_dim,
clip_embeddings_dim=clip_embeddings_dim,
clip_extra_context_tokens=clip_extra_context_tokens
)
else:
self.ipadapter = IPAdapter(
ipadapter,
clip_embeddings,
cross_attention_dim=cross_attention_dim,
clip_embeddings_dim=clip_embeddings_dim,
clip_extra_context_tokens=clip_extra_context_tokens
)
IPA = IPAdapterPlus if self.is_plus else IPAdapter
self.ipadapter = IPA(
ipadapter,
cross_attention_dim=cross_attention_dim,
clip_embeddings_dim=clip_embeddings_dim,
clip_extra_context_tokens=clip_extra_context_tokens
)
self.ipadapter.to(self.device, dtype=self.dtype)
self.image_prompt_embeds, self.uncond_image_prompt_embeds = self.ipadapter.get_image_embeds(clip_embeddings)
self.image_prompt_embeds, self.uncond_image_prompt_embeds = self.ipadapter.get_image_embeds(clip_embed.to(self.device, self.dtype), clip_embed_zeroed.to(self.device, self.dtype))
self.image_prompt_embeds = self.image_prompt_embeds.to(self.device, dtype=self.dtype)
self.uncond_image_prompt_embeds = self.uncond_image_prompt_embeds.to(self.device, dtype=self.dtype)
self.cond_uncond_image_emb = None
work_model = model.clone()
@@ -294,46 +315,38 @@ class IPAdapterApply:
return (work_model,)
def patch_forward(self, number):
device = self.device.type
def forward(n, context_attn2, value_attn2, extra_options):
org_dtype = n.dtype
with torch.autocast(device, dtype=self.dtype):
q = n
k = context_attn2
v = value_attn2
b, _, _ = q.shape
if self.cond_uncond_image_emb is None or self.cond_uncond_image_emb.shape[0] != b:
self.cond_uncond_image_emb = torch.cat([self.uncond_image_prompt_embeds.repeat(b//2, 1, 1), self.image_prompt_embeds.repeat(b//2, 1, 1)], dim=0)
#class IPAClipVisionModel(ClipVisionModel):
# from: https://github.com/comfyanonymous/ComfyUI/blob/f88f7f413afbe04b42c4422e9deedbaa3269ce76/comfy/clip_vision.py#L39
def encode_image_IPAdapter(self, image):
img = torch.clip((255. * image), 0, 255).round().int()
img = list(map(lambda a: a, img))
inputs = self.processor(images=img, return_tensors="pt")
comfy.model_management.load_model_gpu(self.patcher)
pixel_values = inputs['pixel_values'].to(self.load_device)
# k, v for ip_adapter
ip_k = self.ipadapter.ip_layers.to_kvs[number*2](self.cond_uncond_image_emb)
ip_v = self.ipadapter.ip_layers.to_kvs[number*2+1](self.cond_uncond_image_emb)
if self.dtype != torch.float32:
precision_scope = torch.autocast
else:
precision_scope = lambda a, b: contextlib.nullcontext(a)
q, k, v, ip_k, ip_v = map(
lambda t: t.view(b, -1, extra_options["n_heads"], extra_options["dim_head"]).transpose(1, 2),
(q, k, v, ip_k, ip_v),
)
with precision_scope(comfy.model_management.get_autocast_device(self.load_device), torch.float32):
outputs = self.model(pixel_values=pixel_values, output_hidden_states=True)
outputs_zeroed = self.model(torch.zeros_like(pixel_values), output_hidden_states=True).hidden_states[-2]
out = nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
out = out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"])
return outputs, outputs_zeroed
# output of ip_adapter
ip_out = nn.functional.scaled_dot_product_attention(q, ip_k, ip_v, attn_mask=None, dropout_p=0.0, is_causal=False)
ip_out = ip_out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"])
out = out + ip_out * self.weight
return out.to(dtype=org_dtype)
return forward
# Bad developer, bad!
ClipVisionModel.encode_image_IPAdapter = encode_image_IPAdapter
NODE_CLASS_MAPPINGS = {
"IPAdapterModelLoader": IPAdapterModelLoader,
"IPAdapterApply": IPAdapterApply,
"IPAdapterCLIPVisionEncode": IPAdapterCLIPVisionEncode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"IPAdapterModelLoader": "Load IPAdapter Model",
"IPAdapterApply": "Apply IPAdapter",
"IPAdapterCLIPVisionEncode": "CLIP Vision Encode (IPAdapter)"
}
+616
View File
@@ -0,0 +1,616 @@
{
"last_node_id": 33,
"last_link_id": 70,
"nodes": [
{
"id": 8,
"type": "CLIPTextEncode",
"pos": [
790,
310
],
"size": {
"0": 210,
"1": 76
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 6
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
20
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"a girl with fox ears"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 9,
"type": "CLIPTextEncode",
"pos": [
790,
430
],
"size": {
"0": 210,
"1": 76
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 7
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
21
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"horror"
],
"color": "#322",
"bgcolor": "#533"
},
{
"id": 12,
"type": "VAEDecode",
"pos": [
1430,
320
],
"size": {
"0": 140,
"1": 50
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 12
},
{
"name": "vae",
"type": "VAE",
"link": 13
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
31
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 6,
"type": "LoadImage",
"pos": [
120,
200
],
"size": {
"0": 220,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
70
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"example.png",
"image"
]
},
{
"id": 10,
"type": "KSampler",
"pos": [
1070,
320
],
"size": {
"0": 315,
"1": 262
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 44
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 20
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 21
},
{
"name": "latent_image",
"type": "LATENT",
"link": 11,
"slot_index": 3
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
12
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
12,
"fixed",
20,
7,
"dpmpp_2m",
"karras",
1
]
},
{
"id": 31,
"type": "IPAdapterApply",
"pos": [
757,
120
],
"size": {
"0": 315,
"1": 118
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 43
},
{
"name": "ipadapter",
"type": "IPADAPTER",
"link": 45
},
{
"name": "clip_vision_output",
"type": "CLIP_VISION_OUTPUT",
"link": 68
},
{
"name": "mask",
"type": "MASK",
"link": null
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
44
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "IPAdapterApply"
},
"widgets_values": [
1
]
},
{
"id": 20,
"type": "SaveImage",
"pos": [
1430,
420
],
"size": [
526.6748263183599,
579.4905033691409
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 31
}
],
"properties": {},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 1,
"type": "CLIPVisionLoader",
"pos": [
30,
90
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "CLIP_VISION",
"type": "CLIP_VISION",
"links": [
69
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPVisionLoader"
},
"widgets_values": [
"IP-Adapter_sd15_pytorch_model.bin"
]
},
{
"id": 3,
"type": "IPAdapterModelLoader",
"pos": [
430,
60
],
"size": {
"0": 270,
"1": 60
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "IPADAPTER",
"type": "IPADAPTER",
"links": [
45
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "IPAdapterModelLoader"
},
"widgets_values": [
"ip-adapter-plus_sd15.bin"
]
},
{
"id": 33,
"type": "IPAdapterCLIPVisionEncode",
"pos": [
443,
178
],
"size": {
"0": 253.60000610351562,
"1": 50
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip_vision",
"type": "CLIP_VISION",
"link": 69
},
{
"name": "image",
"type": "IMAGE",
"link": 70
}
],
"outputs": [
{
"name": "CLIP_VISION_OUTPUT",
"type": "CLIP_VISION_OUTPUT",
"links": [
68
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "IPAdapterCLIPVisionEncode"
}
},
{
"id": 11,
"type": "EmptyLatentImage",
"pos": [
790,
560
],
"size": {
"0": 210,
"1": 110
},
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
11
],
"shape": 3
}
],
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
512,
512,
1
]
},
{
"id": 7,
"type": "CheckpointLoaderSimple",
"pos": [
387,
295
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 4,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
43
],
"shape": 3,
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
6,
7
],
"shape": 3,
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
13
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"v1-5-pruned-emaonly.safetensors"
]
}
],
"links": [
[
6,
7,
1,
8,
0,
"CLIP"
],
[
7,
7,
1,
9,
0,
"CLIP"
],
[
11,
11,
0,
10,
3,
"LATENT"
],
[
12,
10,
0,
12,
0,
"LATENT"
],
[
13,
7,
2,
12,
1,
"VAE"
],
[
20,
8,
0,
10,
1,
"CONDITIONING"
],
[
21,
9,
0,
10,
2,
"CONDITIONING"
],
[
31,
12,
0,
20,
0,
"IMAGE"
],
[
43,
7,
0,
31,
0,
"MODEL"
],
[
44,
31,
0,
10,
0,
"MODEL"
],
[
45,
3,
0,
31,
1,
"IPADAPTER"
],
[
68,
33,
0,
31,
2,
"CLIP_VISION_OUTPUT"
],
[
69,
1,
0,
33,
0,
"CLIP_VISION"
],
[
70,
6,
0,
33,
1,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
Binary file not shown.

Before

Width:  |  Height:  |  Size: 272 KiB

After

Width:  |  Height:  |  Size: 266 KiB