Merge pull request #74 from cubiq/full

support full model, fix #73
This commit is contained in:
Matteo Spinelli
2023-11-11 08:44:25 +01:00
committed by GitHub
+38 -13
View File
@@ -24,6 +24,22 @@ SD_XL_CHANNELS = [640] * 8 + [1280] * 40 + [1280] * 60 + [640] * 12 + [1280] * 2
def get_filename_list(path):
return [f for f in os.listdir(path) if f.endswith('.bin') or f.endswith('.safetensors')]
class MLPProjModel(torch.nn.Module):
"""SD model with image prompt"""
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024):
super().__init__()
self.proj = torch.nn.Sequential(
torch.nn.Linear(clip_embeddings_dim, clip_embeddings_dim),
torch.nn.GELU(),
torch.nn.Linear(clip_embeddings_dim, cross_attention_dim),
torch.nn.LayerNorm(cross_attention_dim)
)
def forward(self, image_embeds):
clip_extra_context_tokens = self.proj(image_embeds)
return clip_extra_context_tokens
class ImageProjModel(nn.Module):
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
super().__init__()
@@ -148,7 +164,7 @@ def contrast_adaptive_sharpening(image, amount):
return (output)
class IPAdapter(nn.Module):
def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False):
def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False, is_full=False):
super().__init__()
self.clip_embeddings_dim = clip_embeddings_dim
@@ -156,6 +172,7 @@ class IPAdapter(nn.Module):
self.output_cross_attention_dim = output_cross_attention_dim
self.clip_extra_context_tokens = clip_extra_context_tokens
self.is_sdxl = is_sdxl
self.is_full = is_full
self.image_proj_model = self.init_proj() if not is_plus else self.init_proj_plus()
self.image_proj_model.load_state_dict(ipadapter_model["image_proj"])
@@ -171,16 +188,22 @@ class IPAdapter(nn.Module):
return image_proj_model
def init_proj_plus(self):
image_proj_model = Resampler(
dim=self.cross_attention_dim,
depth=4,
dim_head=64,
heads=20 if self.is_sdxl else 12,
num_queries=self.clip_extra_context_tokens,
embedding_dim=self.clip_embeddings_dim,
output_dim=self.output_cross_attention_dim,
ff_mult=4
)
if self.is_full:
image_proj_model = MLPProjModel(
cross_attention_dim=self.cross_attention_dim,
clip_embeddings_dim=self.clip_embeddings_dim
)
else:
image_proj_model = Resampler(
dim=self.cross_attention_dim,
depth=4,
dim_head=64,
heads=20 if self.is_sdxl else 12,
num_queries=self.clip_extra_context_tokens,
embedding_dim=self.clip_embeddings_dim,
output_dim=self.output_cross_attention_dim,
ff_mult=4
)
return image_proj_model
@torch.inference_mode()
@@ -331,7 +354,8 @@ class IPAdapterApply:
self.dtype = model.model.diffusion_model.dtype
self.device = comfy.model_management.get_torch_device()
self.weight = weight
self.is_plus = "latents" in ipadapter["image_proj"]
self.is_full = "proj.0.weight" in ipadapter["image_proj"]
self.is_plus = self.is_full or "latents" in ipadapter["image_proj"]
output_cross_attention_dim = ipadapter["ip_adapter"]["1.to_k_ip.weight"].shape[1]
self.is_sdxl = output_cross_attention_dim == 2048
@@ -371,7 +395,8 @@ class IPAdapterApply:
clip_embeddings_dim=clip_embeddings_dim,
clip_extra_context_tokens=clip_extra_context_tokens,
is_sdxl=self.is_sdxl,
is_plus=self.is_plus
is_plus=self.is_plus,
is_full=self.is_full,
)
self.ipadapter.to(self.device, dtype=self.dtype)