@@ -744,8 +744,6 @@ def prepare_latents_inpaint(
744744 timestep = None ,
745745 is_strength_max = True ,
746746 add_noise = True ,
747- return_noise = False ,
748- return_image_latents = False ,
749747 ):
750748 shape = (
751749 batch_size ,
@@ -768,7 +766,7 @@ def prepare_latents_inpaint(
768766 if image .shape [1 ] == 4 :
769767 image_latents = image .to (device = device , dtype = dtype )
770768 image_latents = image_latents .repeat (batch_size // image_latents .shape [0 ], 1 , 1 , 1 )
771- elif return_image_latents or ( latents is None and not is_strength_max ) :
769+ elif latents is None and not is_strength_max :
772770 image = image .to (device = device , dtype = dtype )
773771 image_latents = self ._encode_vae_image (components , image = image , generator = generator )
774772 image_latents = image_latents .repeat (batch_size // image_latents .shape [0 ], 1 , 1 , 1 )
@@ -786,13 +784,7 @@ def prepare_latents_inpaint(
786784 noise = randn_tensor (shape , generator = generator , device = device , dtype = dtype )
787785 latents = image_latents .to (device )
788786
789- outputs = (latents ,)
790-
791- if return_noise :
792- outputs += (noise ,)
793-
794- if return_image_latents :
795- outputs += (image_latents ,)
787+ outputs = (latents , noise , image_latents )
796788
797789 return outputs
798790
@@ -864,7 +856,7 @@ def __call__(self, components: StableDiffusionXLModularPipeline, state: Pipeline
864856 block_state .height = block_state .image_latents .shape [- 2 ] * components .vae_scale_factor
865857 block_state .width = block_state .image_latents .shape [- 1 ] * components .vae_scale_factor
866858
867- block_state .latents , block_state .noise = self .prepare_latents_inpaint (
859+ block_state .latents , block_state .noise , block_state . image_latents = self .prepare_latents_inpaint (
868860 components ,
869861 block_state .batch_size * block_state .num_images_per_prompt ,
870862 components .num_channels_latents ,
@@ -878,8 +870,6 @@ def __call__(self, components: StableDiffusionXLModularPipeline, state: Pipeline
878870 timestep = block_state .latent_timestep ,
879871 is_strength_max = block_state .is_strength_max ,
880872 add_noise = block_state .add_noise ,
881- return_noise = True ,
882- return_image_latents = False ,
883873 )
884874
885875 # 7. Prepare mask latent variables
0 commit comments