fix(diffusers/pipelines): fix attr bugs in SDi2i pipeline

mindspore-lab · Nov 11, 2024 · a09c25c · a09c25c
1 parent 535d05e
commit a09c25c
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/mindone/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/mindone/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
@@ -510,17 +510,15 @@ def encode_image(self, image, num_images_per_prompt, output_hidden_states=None):
 
         image = image.to(dtype=dtype)
         if output_hidden_states:
-            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True)[2][-2]
             image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
-            uncond_image_enc_hidden_states = self.image_encoder(
-                ops.zeros_like(image), output_hidden_states=True
-            ).hidden_states[-2]
+            uncond_image_enc_hidden_states = self.image_encoder(ops.zeros_like(image), output_hidden_states=True)[2][-2]
             uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
                 num_images_per_prompt, dim=0
             )
             return image_enc_hidden_states, uncond_image_enc_hidden_states
         else:
-            image_embeds = self.image_encoder(image).image_embeds
+            image_embeds = self.image_encoder(image)[0]
             image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
             uncond_image_embeds = ops.zeros_like(image_embeds)