Update pipeline_stable_diffusion_xl_instantid.py

pipeline_stable_diffusion_xl_instantid.py

@@ -217,41 +217,29 @@ class StableDiffusionXLInstantIDPipeline(StableDiffusionXLControlNetPipeline):
             if isinstance(attn_processor, IPAttnProcessor):
                 attn_processor.scale = scale
 
-    def _encode_prompt_image_emb(
-        self,
-        image_embeds,
-        num_images_per_prompt,
-        device,
-        do_classifier_free_guidance,
-        prompt_image_emb=None,   # 5ème argument optionnel
-    ):
-        # Cas où on ne fournit pas d'image_embeds : on désactive ce chemin
-        if image_embeds is None:
-            return None, None
-
-        # Dans ta version d’origine, ils assignent prompt_image_emb depuis image_embeds
-        if prompt_image_emb is None:
-            prompt_image_emb = image_embeds
-
-        # S'assurer que c'est un tensor
-        if not isinstance(prompt_image_emb, torch.Tensor):
-            prompt_image_emb = torch.tensor(prompt_image_emb, device=device)
+    def _encode_prompt_image_emb(self, prompt_image_emb, device, num_images_per_prompt, dtype, do_classifier_free_guidance):
+        
+        if isinstance(prompt_image_emb, torch.Tensor):
+            prompt_image_emb = prompt_image_emb.clone().detach()
         else:
-            prompt_image_emb = prompt_image_emb.to(device)
-
-        # Répéter pour num_images_per_prompt
-        bs = prompt_image_emb.shape[0]
-        if bs != num_images_per_prompt:
-            prompt_image_emb = prompt_image_emb.repeat_interleave(num_images_per_prompt, dim=0)
-
-        # Classifier-free guidance : on crée un embedding négatif si besoin
+            prompt_image_emb = torch.tensor(prompt_image_emb)
+            
+        prompt_image_emb = prompt_image_emb.reshape([1, -1, self.image_proj_model_in_features])
+        
         if do_classifier_free_guidance:
-            negative_prompt_image_emb = torch.zeros_like(prompt_image_emb, device=device)
-            prompt_image_emb = torch.cat([negative_prompt_image_emb, prompt_image_emb], dim=0)
+            prompt_image_emb = torch.cat([torch.zeros_like(prompt_image_emb), prompt_image_emb], dim=0)
         else:
-            negative_prompt_image_emb = None
+            prompt_image_emb = torch.cat([prompt_image_emb], dim=0)
+        
+        prompt_image_emb = prompt_image_emb.to(device=self.image_proj_model.latents.device, 
+                                               dtype=self.image_proj_model.latents.dtype)
+        prompt_image_emb = self.image_proj_model(prompt_image_emb)
 
-        return prompt_image_emb, negative_prompt_image_emb
+        bs_embed, seq_len, _ = prompt_image_emb.shape
+        prompt_image_emb = prompt_image_emb.repeat(1, num_images_per_prompt, 1)
+        prompt_image_emb = prompt_image_emb.view(bs_embed * num_images_per_prompt, seq_len, -1)
+        
+        return prompt_image_emb.to(device=device, dtype=dtype)
 
     @torch.no_grad()
     @replace_example_docstring(EXAMPLE_DOC_STRING)