Fix/Enable all schedulers for in-painting (huggingface#1331)

* inpaint fix k lms

* onnox as well

* up

Author: patrickvonplaten

pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py

@@ -408,8 +408,8 @@ def __call__(
             # expand the latents if we are doing classifier free guidance
             latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
             # concat latents, mask, masked_image_latnets in the channel dimension
-            latent_model_input = np.concatenate([latent_model_input, mask, masked_image_latents], axis=1)
             latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)
+            latent_model_input = np.concatenate([latent_model_input, mask, masked_image_latents], axis=1)
             latent_model_input = latent_model_input.cpu().numpy()
 
             # predict the noise residual

pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -586,9 +586,8 @@ def __call__(
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
 
             # concat latents, mask, masked_image_latents in the channel dimension
-            latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
-
             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
 
             # predict the noise residual
             noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample