commit fb95e15653c1e70cacbe3a5cccc0f40de1307d61 Author: Raphael Walker Date: Wed Dec 4 15:21:00 2024 +0100 initial commit diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..f51407b --- /dev/null +++ b/README.md @@ -0,0 +1,7 @@ +# ComfyUI-IPAdapter-SD3 + +ComfyUI implementation of the [InstantX IP-Adapter for SD3.5 Large](https://huggingface.co/InstantX/SD3.5-Large-IP-Adapter). + +## Installation + +Download [`ip-adapter.bin` from the original repository](https://huggingface.co/InstantX/SD3.5-Large-IP-Adapter/blob/main/ip-adapter.bin), and place it in the `models/ipadapter` folder of your ComfyUI installation. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/models/jointblock.py b/models/jointblock.py new file mode 100644 index 0000000..4c361df --- /dev/null +++ b/models/jointblock.py @@ -0,0 +1,217 @@ +from typing import Optional + +import torch +from torch import nn +from torch.nn import functional as F +from einops import rearrange + +from comfy.ldm.modules.attention import optimized_attention +from comfy.ldm.modules.diffusionmodules.mmdit import ( + DismantledBlock, + RMSNorm, + JointBlock, +) + + +class AdaLayerNorm(nn.Module): + """ + Norm layer adaptive layer norm zero (adaLN-Zero). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the embeddings dictionary. + """ + + def __init__(self, embedding_dim: int, time_embedding_dim=None, mode="normal"): + super().__init__() + + self.silu = nn.SiLU() + num_params_dict = dict( + zero=6, + normal=2, + ) + num_params = num_params_dict[mode] + self.linear = nn.Linear( + time_embedding_dim or embedding_dim, num_params * embedding_dim, bias=True + ) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) + self.mode = mode + + def forward( + self, + x, + hidden_dtype=None, + emb=None, + ): + emb = self.linear(self.silu(emb)) + if self.mode == "normal": + shift_msa, scale_msa = emb.chunk(2, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x + + elif self.mode == "zero": + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk( + 6, dim=1 + ) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class IPAttnProcessor(nn.Module): + + def __init__( + self, + hidden_size=None, + cross_attention_dim=None, + ip_hidden_states_dim=None, + ip_encoder_hidden_states_dim=None, + head_dim=None, + timesteps_emb_dim=1280, + ): + super().__init__() + + self.norm_ip = AdaLayerNorm( + ip_hidden_states_dim, time_embedding_dim=timesteps_emb_dim + ) + self.to_k_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False) + self.to_v_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False) + self.norm_q = RMSNorm(head_dim, 1e-6) + self.norm_k = RMSNorm(head_dim, 1e-6) + self.norm_ip_k = RMSNorm(head_dim, 1e-6) + + def forward( + self, + ip_hidden_states, + img_query, + img_key=None, + img_value=None, + t_emb=None, + n_heads=1, + ): + if ip_hidden_states is None: + return None + + if not hasattr(self, "to_k_ip") or not hasattr(self, "to_v_ip"): + return None + + # norm ip input + norm_ip_hidden_states = self.norm_ip(ip_hidden_states, emb=t_emb) + + # to k and v + ip_key = self.to_k_ip(norm_ip_hidden_states) + ip_value = self.to_v_ip(norm_ip_hidden_states) + + # reshape + img_query = rearrange(img_query, "b l (h d) -> b h l d", h=n_heads) + img_key = rearrange(img_key, "b l (h d) -> b h l d", h=n_heads) + # note that the image is in a different shape: b l h d + # so we transpose to b h l d + # or do we have to transpose here? + img_value = torch.transpose(img_value, 1, 2) + ip_key = rearrange(ip_key, "b l (h d) -> b h l d", h=n_heads) + ip_value = rearrange(ip_value, "b l (h d) -> b h l d", h=n_heads) + + # norm + img_query = self.norm_q(img_query) + img_key = self.norm_k(img_key) + ip_key = self.norm_ip_k(ip_key) + + # cat img + key = torch.cat([img_key, ip_key], dim=2) + value = torch.cat([img_value, ip_value], dim=2) + + # + ip_hidden_states = F.scaled_dot_product_attention( + img_query, key, value, dropout_p=0.0, is_causal=False + ) + ip_hidden_states = rearrange(ip_hidden_states, "b h l d -> b l (h d)") + ip_hidden_states = ip_hidden_states.to(img_query.dtype) + return ip_hidden_states + + +class JointBlockIPWrapper: + """To be used as a patch_replace with Comfy""" + + def __init__( + self, + original_block: JointBlock, + adapter: IPAttnProcessor, + ip_options=None, + ): + self.original_block = original_block + self.adapter = adapter + self.ip_options = ip_options + + def block_mixing(self, context, x, context_block, x_block, c): + """ + Comes from mmdit.py. Modified to add ipadapter attention. + """ + context_qkv, context_intermediates = context_block.pre_attention(context, c) + + if x_block.x_block_self_attn: + x_qkv, x_qkv2, x_intermediates = x_block.pre_attention_x(x, c) + else: + x_qkv, x_intermediates = x_block.pre_attention(x, c) + + qkv = tuple(torch.cat((context_qkv[j], x_qkv[j]), dim=1) for j in range(3)) + + attn = optimized_attention( + qkv[0], + qkv[1], + qkv[2], + heads=x_block.attn.num_heads, + ) + context_attn, x_attn = ( + attn[:, : context_qkv[0].shape[1]], + attn[:, context_qkv[0].shape[1] :], + ) + # IP-Adapter + ip_attn = self.adapter( + self.ip_options["hidden_states"], + *x_qkv, + self.ip_options["t_emb"], + x_block.attn.num_heads, + ) + x_attn = x_attn + ip_attn * self.ip_options["weight"] + + # Everything else is unchanged + if not context_block.pre_only: + context = context_block.post_attention(context_attn, *context_intermediates) + + else: + context = None + if x_block.x_block_self_attn: + attn2 = optimized_attention( + x_qkv2[0], + x_qkv2[1], + x_qkv2[2], + heads=x_block.attn2.num_heads, + ) + x = x_block.post_attention_x(x_attn, attn2, *x_intermediates) + else: + x = x_block.post_attention(x_attn, *x_intermediates) + return context, x + + def __call__(self, args, _): + # Code from mmdit.py: + # in this case, we're blocks_replace[("double_block", i)] + # note that although we're passed the original block, + # we can't actually get it from inside its wrapper + # (which would simplify the whole code...) + # ``` + # def block_wrap(args): + # out = {} + # out["txt"], out["img"] = self.joint_blocks[i](args["txt"], args["img"], c=args["vec"]) + # return out + # out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": c_mod}, {"original_block": block_wrap}) + # context = out["txt"] + # x = out["img"] + # ``` + c, x = self.block_mixing( + args["txt"], + args["img"], + self.original_block.context_block, + self.original_block.x_block, + c=args["vec"], + ) + return {"txt": c, "img": x} diff --git a/models/resampler.py b/models/resampler.py new file mode 100644 index 0000000..f058669 --- /dev/null +++ b/models/resampler.py @@ -0,0 +1,304 @@ +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +import math + +import torch +import torch.nn as nn + +from diffusers.models.embeddings import Timesteps, TimestepEmbedding + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the + embeddings. :return: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + # (bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents, shift=None, scale=None): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + if shift is not None and scale is not None: + latents = latents * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose( + -2, -1 + ) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + *args, + **kwargs, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) + + +class TimeResampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + timestep_in_dim=320, + timestep_flip_sin_to_cos=True, + timestep_freq_shift=0, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + # msa + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + # ff + FeedForward(dim=dim, mult=ff_mult), + # adaLN + nn.Sequential(nn.SiLU(), nn.Linear(dim, 4 * dim, bias=True)), + ] + ) + ) + + # time + self.time_proj = Timesteps( + timestep_in_dim, timestep_flip_sin_to_cos, timestep_freq_shift + ) + self.time_embedding = TimestepEmbedding(timestep_in_dim, dim, act_fn="silu") + + # adaLN + # self.adaLN_modulation = nn.Sequential( + # nn.SiLU(), + # nn.Linear(timestep_out_dim, 6 * timestep_out_dim, bias=True) + # ) + + def forward(self, x, timestep, need_temb=False): + timestep_emb = self.embedding_time(x, timestep) # bs, dim + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + x = x + timestep_emb[:, None] + + for attn, ff, adaLN_modulation in self.layers: + shift_msa, scale_msa, shift_mlp, scale_mlp = adaLN_modulation( + timestep_emb + ).chunk(4, dim=1) + latents = attn(x, latents, shift_msa, scale_msa) + latents + + res = latents + for idx_ff in range(len(ff)): + layer_ff = ff[idx_ff] + latents = layer_ff(latents) + if idx_ff == 0 and isinstance(layer_ff, nn.LayerNorm): # adaLN + latents = latents * ( + 1 + scale_mlp.unsqueeze(1) + ) + shift_mlp.unsqueeze(1) + latents = latents + res + + # latents = ff(latents) + latents + + latents = self.proj_out(latents) + latents = self.norm_out(latents) + + if need_temb: + return latents, timestep_emb + else: + return latents + + def embedding_time(self, sample, timestep): + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # timesteps does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, None) + return emb + + +if __name__ == "__main__": + model = TimeResampler( + dim=1280, + depth=4, + dim_head=64, + heads=20, + num_queries=16, + embedding_dim=512, + output_dim=2048, + ff_mult=4, + timestep_in_dim=320, + timestep_flip_sin_to_cos=True, + timestep_freq_shift=0, + in_channel_extra_emb=2048, + ) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..e275d9d --- /dev/null +++ b/nodes.py @@ -0,0 +1,191 @@ +import os +import logging +from functools import partial + +import torch +import folder_paths +from transformers import AutoProcessor, SiglipVisionModel +from PIL import Image +import numpy as np + +from .models.resampler import TimeResampler +from .models.jointblock import JointBlockIPWrapper, IPAttnProcessor + +MODELS_DIR = os.path.join(folder_paths.models_dir, "ipadapter") +if "ipadapter" not in folder_paths.folder_names_and_paths: + current_paths = [MODELS_DIR] +else: + current_paths, _ = folder_paths.folder_names_and_paths["ipadapter"] +folder_paths.folder_names_and_paths["ipadapter"] = ( + current_paths, + folder_paths.supported_pt_extensions, +) + + +def patch(patcher, ip_procs, curried_resampler, weight=1.0): + """ + Patches a model_sampler to add the ipadapter + """ + mmdit = patcher.model.diffusion_model + timestep_schedule_max = patcher.model.model_config.sampling_settings.get("timesteps", 1000) + # hook the model's forward function + # so that when it gets called, we can grab the timestep and send it to the resampler + ip_options = {"hidden_states": None, "t_emb": None, "t": None, "weight": weight} + + def ddit_wrapper(forward, args): + # these will need to be accessible to the IPAdapters + # this is between 0 and 1, so the adapters can calculate start_point and end_point + ip_options["t"] = args["timestep"] + # but the resampler wants between 0 and MAX_STEPS + timestep_expanded = args["timestep"] * timestep_schedule_max + image_emb, t_emb = curried_resampler(timestep_expanded) + ip_options["hidden_states"] = image_emb + ip_options["t_emb"] = t_emb + # is this the right way to call? + return forward(args["input"], args["timestep"], **args["c"]) + + patcher.set_model_unet_function_wrapper(ddit_wrapper) + # patch each dit block + for i, block in enumerate(mmdit.joint_blocks): + wrapper = JointBlockIPWrapper(block, ip_procs[i], ip_options) + patcher.set_model_patch_replace(wrapper, "dit", "double_block", i) + + +class SD3IPAdapter: + def __init__(self, encoder_path: str, checkpoint: str, device): + self.device = device + # load image encoder + self.encoder = SiglipVisionModel.from_pretrained(encoder_path).to( + self.device, dtype=torch.float16 + ) + self.processor = AutoProcessor.from_pretrained(encoder_path) + # load the checkpoint right away + self.state_dict = torch.load( + os.path.join(MODELS_DIR, checkpoint), + map_location=self.device, + weights_only=True, + ) + # todo: infer some of the params from the checkpoint instead of hardcoded + self.resampler = TimeResampler( + dim=1280, + depth=4, + dim_head=64, + heads=20, + num_queries=64, + embedding_dim=1152, + output_dim=2432, + ff_mult=4, + timestep_in_dim=320, + timestep_flip_sin_to_cos=True, + timestep_freq_shift=0, + ) + self.resampler.eval() + self.resampler.to(self.device, dtype=torch.float16) + self.resampler.load_state_dict(self.state_dict["image_proj"]) + + # now we'll create the attention processors + # ip_adapter.keys looks like [0.proj, 0.to_k, ..., 1.proj, 1.to_k, ...] + n_procs = len( + set(x.split(".")[0] for x in self.state_dict["ip_adapter"].keys()) + ) + self.procs = torch.nn.ModuleList( + [ + # this is hardcoded for SD3.5L + IPAttnProcessor( + hidden_size=2432, + cross_attention_dim=2432, + ip_hidden_states_dim=2432, + ip_encoder_hidden_states_dim=2432, + head_dim=64, + timesteps_emb_dim=1280, + ).to(self.device, dtype=torch.float16) + for _ in range(n_procs) + ] + ) + self.procs.load_state_dict(self.state_dict["ip_adapter"]) + + @torch.inference_mode() + def encode(self, image): + clip_image = self.processor.image_processor( + image, return_tensors="pt", do_rescale=False + ).pixel_values + clip_image_embeds = self.encoder( + clip_image.to(self.device, dtype=self.encoder.dtype), + output_hidden_states=True, + ).hidden_states[-2] + clip_image_embeds = torch.cat( + [torch.zeros_like(clip_image_embeds), clip_image_embeds], dim=0 + ) + clip_image_embeds = clip_image_embeds.to(dtype=torch.float16) + return clip_image_embeds + + +class IPAdapterSD3Loader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ipadapter": (folder_paths.get_filename_list("ipadapter"),), + "clip_vision": (["google/siglip-so400m-patch14-384"],), + "provider": (["cuda", "cpu", "mps"],), + } + } + + RETURN_TYPES = ("IP_ADAPTER_SD3_INSTANTX",) + RETURN_NAMES = ("ipadapter",) + FUNCTION = "load_model" + CATEGORY = "InstantXNodes" + + def load_model(self, ipadapter, clip_vision, provider): + logging.info("Loading InstantX IPAdapter SD3 model.") + model = SD3IPAdapter(clip_vision, ipadapter, provider) + return (model,) + + +class ApplyIPAdapterSD3: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "ipadapter": ("IP_ADAPTER_SD3_INSTANTX",), + "image": ("IMAGE",), + "weight": ( + "FLOAT", + {"default": 1.0, "min": -1.0, "max": 5.0, "step": 0.05}, + ), + "start_percent": ( + "FLOAT", + {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + "end_percent": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply_ipadapter" + CATEGORY = "InstantXNodes" + + def apply_ipadapter( + self, model, ipadapter, image, weight, start_percent, end_percent + ): + # set model + new_model = model.clone() + embeds = ipadapter.encode(image) + curried_resampler = partial(ipadapter.resampler, embeds, need_temb=True) + patch(new_model, ipadapter.procs, curried_resampler, weight=weight) + return (new_model,) + + +NODE_CLASS_MAPPINGS = { + "IPAdapterSD3Loader": IPAdapterSD3Loader, + "ApplyIPAdapterSD3": ApplyIPAdapterSD3, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IPAdapterSD3Loader": "Load IPAdapter SD3 Model", + "ApplyIPAdapterSD3": "Apply IPAdapter SD3 Model", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5a66498 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +torch +pillow +einops==0.8.0 +transformers>=4.37.2 +diffusers +sentencepiece>=0.2.0 +protobuf>=4.25.5