Update ip_adapter/ip_adapter.py

ip_adapter/ip_adapter.py

@@ -30,51 +30,17 @@ class ImageProjModel(torch.nn.Module):
     def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
         super().__init__()
 
-        # cross_attention_dim = 768
-        # clip_extra_context_tokens = 4
-        # clip_embeddings_dim = 1024
         self.cross_attention_dim = cross_attention_dim
         self.clip_extra_context_tokens = clip_extra_context_tokens
-        # 创建了一个线性层self.proj，将clip_embeddings_dim作为输入维度，将self.clip_extra_context_tokens * cross_attention_dim作为输出维度。
         self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
-        # self.proj_1 = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
-        #
-        # # 访问线性层的权重参数
-        # weights = self.proj.weight
-        # print("proj_weights")
-        # print(weights)
-        # # 访问线性层的权重参数
-        # weights_1 = self.proj_1.weight
-        # print("proj_1_weights")
-        # print(weights_1)
-        #
-        # # 访问线性层的偏置参数
-        # bias = self.proj.bias
-        # print("proj_bias")
-        # print(bias)
-        # # 访问线性层的偏置参数
-        # bias_1 = self.proj_1.bias
-        # print("proj_1_bias")
-        # print(bias_1)
-
-
-        # 接着，它创建了一个LayerNorm层self.norm，将cross_attention_dim作为输入维度
-        # LayerNorm层能对每个通道进行归一化处理，确保每个通道均值方差一致，使得每个通道的特征分布相对一致，帮助模型学习特征
         self.norm = torch.nn.LayerNorm(cross_attention_dim)
 
     def forward(self, image_embeds):
-        # 在前向传播函数中，它接受image_embeds作为输入，然后将其赋值给embeds。
         embeds = image_embeds
-        # embeds.shape = [1,1024]
-        # self.proj(embeds).shape = [1,3072]
-        # 接着，它使用self.proj对embeds进行线性变换，并将结果reshape
         clip_extra_context_tokens = self.proj(embeds).reshape(
             -1, self.clip_extra_context_tokens, self.cross_attention_dim
         )
-        # clip_extra_context_tokens.shape = [1,4,768]
-        # 然后，它将结果传入self.norm进行LayerNorm操作，并返回处理后的结果clip_extra_context_tokens。
         clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
-        # clip_extra_context_tokens.shape = [1,4,768]
         return clip_extra_context_tokens
 
 
@@ -110,7 +76,7 @@ class IPAdapter:
 
         # load image encoder
         self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to(
-            self.device, dtype=torch.float16
+            self.device, dtype=torch.float32
         )
         self.clip_image_processor = CLIPImageProcessor()
         # image proj model
@@ -123,20 +89,14 @@ class IPAdapter:
             cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
             clip_embeddings_dim=self.image_encoder.config.projection_dim,
             clip_extra_context_tokens=self.num_tokens,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=torch.float32)
         return image_proj_model
 
     def set_ip_adapter(self):
-        # 首先，它获取了self.pipe.unet中的unet，
         unet = self.pipe.unet
-        # 并初始化了一个空的字典attn_procs
         attn_procs = {}
-        # 然后，它遍历unet.attn_processors中的每个键名name
         for name in unet.attn_processors.keys():
-            # 在循环中，它根据name的不同情况设置cross_attention_dim和hidden_size
-            # 如果name以"attn1.processor"结尾，那么cross_attention_dim被设置为None；否则，它被设置为unet.config.cross_attention_dim。
             cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
-            # 接着，根据name的前缀不同，设置了hidden_size的值
             if name.startswith("mid_block"):
                 hidden_size = unet.config.block_out_channels[-1]
             elif name.startswith("up_blocks"):
@@ -145,12 +105,9 @@ class IPAdapter:
             elif name.startswith("down_blocks"):
                 block_id = int(name[len("down_blocks.")])
                 hidden_size = unet.config.block_out_channels[block_id]
-            # 接下来，根据cross_attention_dim的值，为每个name创建了一个对应的AttnProcessor或IPAttnProcessor，并将其加入attn_procs字典中最后
             if cross_attention_dim is None:
-                #print("initialization：attn_procs[name] = AttnProcessor()")
                 attn_procs[name] = AttnProcessor()
             else:
-                #print("initialization：attn_procs[name] = IPAttnProcessor()")
                 attn_procs[name] = IPAttnProcessor(
                     hidden_size= hidden_size,
                     cross_attention_dim=cross_attention_dim,
@@ -158,9 +115,7 @@ class IPAdapter:
                     num_tokens=self.num_tokens,
                     Control_factor=self.Control_factor,
                     IP_factor=self.IP_factor,
-                ).to(self.device, dtype=torch.float16)
-        # 调用unet.set_attn_processor(attn_procs)来设置unet的注意力处理器
-        # 同时调用self.pipe.controlnet.set_attn_processor(CNAttnProcessor(num_tokens=self.num_tokens))来设置self.pipe.controlnet的注意力处理器。
+                ).to(self.device, dtype=torch.float32)
         unet.set_attn_processor(attn_procs)
         #self.pipe.controlnet.set_attn_processor(CNAttnProcessor(num_tokens=self.num_tokens))
         if hasattr(self.pipe, "controlnet"):
@@ -171,12 +126,8 @@ class IPAdapter:
                 self.pipe.controlnet.set_attn_processor(CNAttnProcessor(num_tokens=self.num_tokens))
 
     def load_ip_adapter(self):
-        # 该方法用于加载IP适配器的状态。然后，它使用safe_open函数打开self.ip_ckpt文件，并遍历文件中的键名。
-        # 首先，它检查self.ip_ckpt的文件扩展名是否为".safetensors"。
         if os.path.splitext(self.ip_ckpt)[-1] == ".safetensors":
-            # 如果是，它创建了一个空的state_dict字典，包含"image_proj"和"ip_adapter"两个键。
             state_dict = {"image_proj": {}, "ip_adapter": {}}
-            # 对于以"image_proj."开头的键名，它将对应的张量存入state_dict["image_proj"]中；对于以"ip_adapter."开头的键名，它将对应的张量存入state_dict["ip_adapter"]中。
             with safe_open(self.ip_ckpt, framework="pt", device="cpu") as f:
                 for key in f.keys():
                     if key.startswith("image_proj."):
@@ -184,12 +135,7 @@ class IPAdapter:
                     elif key.startswith("ip_adapter."):
                         state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
         else:
-            # 如果self.ip_ckpt的文件扩展名不是".safetensors"，那么它直接使用torch.load函数加载self.ip_ckpt文件的状态，并将其存入state_dict中。
             state_dict = torch.load(self.ip_ckpt, map_location="cpu")
-        # 这段代码中的两行分别用于加载预训练模型的参数。
-        # 第一行使用load_state_dict方法将state_dict中的"image_proj"部分加载到self.image_proj_model中
-        # 而第二行则尝试将state_dict中的"ip_adapter"部分加载到ip_layers中。
-        # 需要注意的是，ip_layers是一个ModuleList，它包含了多个attn_processors，因此在尝试加载"ip_adapter"部分时，需要确保state_dict中的键能够与ip_layers中的各个子模块对应上。
         self.image_proj_model.load_state_dict(state_dict["image_proj"])
         ip_layers = torch.nn.ModuleList(self.pipe.unet.attn_processors.values())
         ip_layers.load_state_dict(state_dict["ip_adapter"])
@@ -200,14 +146,9 @@ class IPAdapter:
             if isinstance(pil_image, Image.Image):
                 pil_image = [pil_image]
             clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
-            clip_image_embeds = self.image_encoder(clip_image.to(self.device, dtype=torch.float16)).image_embeds
-
-            # clip_imageBroken = self.clip_image_processor(images=image_broken, return_tensors="pt").pixel_values
-            # clip_imageBroken_embeds = self.image_encoder(clip_imageBroken.to(self.device, dtype=torch.float16)).image_embeds
-            # clip_image_embeds.shape: torch.Size([1, 1024])
-            # style_vector = clip_image_embeds-clip_imageBroken_embeds
+            clip_image_embeds = self.image_encoder(clip_image.to(self.device, dtype=torch.float32)).image_embeds
         else:
-            clip_image_embeds = clip_image_embeds.to(self.device, dtype=torch.float16)
+            clip_image_embeds = clip_image_embeds.to(self.device, dtype=torch.float32)
 
 
         # image_prompt_embeds = self.image_proj_model(style_vector)
@@ -382,7 +323,7 @@ class IPAdapterPlus(IPAdapter):
             embedding_dim=self.image_encoder.config.hidden_size,
             output_dim=self.pipe.unet.config.cross_attention_dim,
             ff_mult=4,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=torch.float32)
         return image_proj_model
 
     @torch.inference_mode()
@@ -390,7 +331,7 @@ class IPAdapterPlus(IPAdapter):
         if isinstance(pil_image, Image.Image):
             pil_image = [pil_image]
         clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
-        clip_image = clip_image.to(self.device, dtype=torch.float16)
+        clip_image = clip_image.to(self.device, dtype=torch.float32)
         clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
         image_prompt_embeds = self.image_proj_model(clip_image_embeds)
         uncond_clip_image_embeds = self.image_encoder(
@@ -408,7 +349,7 @@ class IPAdapterFull(IPAdapterPlus):
         image_proj_model = MLPProjModel(
             cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
             clip_embeddings_dim=self.image_encoder.config.hidden_size,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=torch.float32)
         return image_proj_model
 
 # image_proj_model = Resampler(
@@ -425,7 +366,7 @@ class IPAdapterPlusXL(IPAdapter):
             embedding_dim=self.image_encoder.config.hidden_size,
             output_dim=self.pipe.unet.config.cross_attention_dim,
             ff_mult=4,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=torch.float32)
         return image_proj_model
 
     @torch.inference_mode()
@@ -433,7 +374,7 @@ class IPAdapterPlusXL(IPAdapter):
         if isinstance(pil_image, Image.Image):
             pil_image = [pil_image]
         clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
-        clip_image = clip_image.to(self.device, dtype=torch.float16)
+        clip_image = clip_image.to(self.device, dtype=torch.float32)
         clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
         image_prompt_embeds = self.image_proj_model(clip_image_embeds)
         uncond_clip_image_embeds = self.image_encoder(