| 47 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |
| 48 | """ |
| 49 | class IPAdapter(nn.Module): |
| 50 | 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, is_faceid=False, is_portrait_unnorm=False, is_kwai_kolors=False, encoder_hid_proj=None, weight_kolors=1.0): |
| 51 | super().__init__() |
| 52 | |
| 53 | self.clip_embeddings_dim = clip_embeddings_dim |
| 54 | self.cross_attention_dim = cross_attention_dim |
| 55 | self.output_cross_attention_dim = output_cross_attention_dim |
| 56 | self.clip_extra_context_tokens = clip_extra_context_tokens |
| 57 | self.is_sdxl = is_sdxl |
| 58 | self.is_full = is_full |
| 59 | self.is_plus = is_plus |
| 60 | self.is_portrait_unnorm = is_portrait_unnorm |
| 61 | self.is_kwai_kolors = is_kwai_kolors |
| 62 | |
| 63 | if is_faceid and not is_portrait_unnorm: |
| 64 | self.image_proj_model = self.init_proj_faceid() |
| 65 | elif is_full: |
| 66 | self.image_proj_model = self.init_proj_full() |
| 67 | elif is_plus or is_portrait_unnorm: |
| 68 | self.image_proj_model = self.init_proj_plus() |
| 69 | else: |
| 70 | self.image_proj_model = self.init_proj() |
| 71 | |
| 72 | self.image_proj_model.load_state_dict(ipadapter_model["image_proj"]) |
| 73 | self.ip_layers = To_KV(ipadapter_model["ip_adapter"], encoder_hid_proj=encoder_hid_proj, weight_kolors=weight_kolors) |
| 74 | |
| 75 | self.multigpu_clones = {} |
| 76 | |
| 77 | def create_multigpu_clone(self, device): |
| 78 | if device not in self.multigpu_clones: |
| 79 | orig_multigpu_clones = self.multigpu_clones |
| 80 | try: |
| 81 | self.multigpu_clones = {} |
| 82 | new_clone = copy.deepcopy(self) |
| 83 | new_clone = new_clone.to(device) |
| 84 | orig_multigpu_clones[device] = new_clone |
| 85 | finally: |
| 86 | self.multigpu_clones = orig_multigpu_clones |
| 87 | |
| 88 | def get_multigpu_clone(self, device): |
| 89 | return self.multigpu_clones.get(device, self) |
| 90 | |
| 91 | def init_proj(self): |
| 92 | image_proj_model = ImageProjModel( |
| 93 | cross_attention_dim=self.cross_attention_dim, |
| 94 | clip_embeddings_dim=self.clip_embeddings_dim, |
| 95 | clip_extra_context_tokens=self.clip_extra_context_tokens |
| 96 | ) |
| 97 | return image_proj_model |
| 98 | |
| 99 | def init_proj_plus(self): |
| 100 | image_proj_model = Resampler( |
| 101 | dim=self.cross_attention_dim, |
| 102 | depth=4, |
| 103 | dim_head=64, |
| 104 | heads=20 if self.is_sdxl and not self.is_kwai_kolors else 12, |
| 105 | num_queries=self.clip_extra_context_tokens, |
| 106 | embedding_dim=self.clip_embeddings_dim, |