| 57 | return out_img_pil |
| 58 | |
| 59 | class InstantID(torch.nn.Module): |
| 60 | def __init__(self, instantid_model, cross_attention_dim=1280, output_cross_attention_dim=1024, clip_embeddings_dim=512, clip_extra_context_tokens=16): |
| 61 | super().__init__() |
| 62 | |
| 63 | self.clip_embeddings_dim = clip_embeddings_dim |
| 64 | self.cross_attention_dim = cross_attention_dim |
| 65 | self.output_cross_attention_dim = output_cross_attention_dim |
| 66 | self.clip_extra_context_tokens = clip_extra_context_tokens |
| 67 | |
| 68 | self.image_proj_model = self.init_proj() |
| 69 | |
| 70 | self.image_proj_model.load_state_dict(instantid_model["image_proj"]) |
| 71 | self.ip_layers = To_KV(instantid_model["ip_adapter"]) |
| 72 | |
| 73 | def init_proj(self): |
| 74 | image_proj_model = Resampler( |
| 75 | dim=self.cross_attention_dim, |
| 76 | depth=4, |
| 77 | dim_head=64, |
| 78 | heads=20, |
| 79 | num_queries=self.clip_extra_context_tokens, |
| 80 | embedding_dim=self.clip_embeddings_dim, |
| 81 | output_dim=self.output_cross_attention_dim, |
| 82 | ff_mult=4 |
| 83 | ) |
| 84 | return image_proj_model |
| 85 | |
| 86 | @torch.inference_mode() |
| 87 | def get_image_embeds(self, clip_embed, clip_embed_zeroed): |
| 88 | #image_prompt_embeds = clip_embed.clone().detach() |
| 89 | image_prompt_embeds = self.image_proj_model(clip_embed) |
| 90 | #uncond_image_prompt_embeds = clip_embed_zeroed.clone().detach() |
| 91 | uncond_image_prompt_embeds = self.image_proj_model(clip_embed_zeroed) |
| 92 | |
| 93 | return image_prompt_embeds, uncond_image_prompt_embeds |
| 94 | |
| 95 | class ImageProjModel(torch.nn.Module): |
| 96 | def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): |