StableDiffusionImg2ImgPipeline can't load by onnx revision

[LowinLi]

Describe the bug

when load StableDiffusionImg2ImgPipeline through from_pretrained by revision="onnx". an Exception raised:

File ~/venv/lib/python3.8/site-packages/diffusers/pipeline_utils.py:391, in DiffusionPipeline.from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
    388     init_kwargs[name] = loaded_sub_model  # UNet(...), # DiffusionSchedule(...)
    390 # 4. Instantiate the pipeline
--> 391 model = pipeline_class(**init_kwargs)
    392 return model

TypeError: __init__() missing 1 required positional argument: 'vae'

Reproduction

from diffusers import StableDiffusionImg2ImgPipeline

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4", revision="onnx", use_auth_token=token, provider="CPUExecutionProvider"
)

Logs

TypeError                                 Traceback (most recent call last)
Cell In [9], line 9
      6 token = "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"
      7 # load the pipeline
----> 9 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     10     "CompVis/stable-diffusion-v1-4", revision="onnx", use_auth_token=token, provider="CPUExecutionProvider", force_download=True
     11 )
     13 # let's download an initial image
     14 url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"

File ~/venv/lib/python3.8/site-packages/diffusers/pipeline_utils.py:391, in DiffusionPipeline.from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
    388     init_kwargs[name] = loaded_sub_model  # UNet(...), # DiffusionSchedule(...)
    390 # 4. Instantiate the pipeline
--> 391 model = pipeline_class(**init_kwargs)
    392 return model

TypeError: __init__() missing 1 required positional argument: 'vae'

System Info

• diffusers version: 0.3.0
• Platform: Linux-5.13.0-41-generic-x86_64-with-glibc2.29
• Python version: 3.8.1
• PyTorch version (GPU?): 1.10.0+cu102 (True)
• Huggingface_hub version: 0.9.1
• Transformers version: 4.22.2
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patil-suraj]

Hi, thanks for the issue! The StableDiffusionImg2ImgPipeline only works with PyTorch models, it doesn't support onnx. For onnx you should use StableDiffusionOnnxPipeline, but note that it does not implement imag2image, but should be easy to modify to add imahe2image in it.

We will need to pass init_image as in the StableDiffusionImg2ImgPipeline and prepare the initial latents, think you pretty much copy that part into the onnx pipeline

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Lines 217 to 239 in 7258dc4

	if isinstance(init_image, PIL.Image.Image):
	init_image = preprocess(init_image)
	# encode the init image into latents and scale the latents
	init_latent_dist = self.vae.encode(init_image.to(self.device)).latent_dist
	init_latents = init_latent_dist.sample(generator=generator)
	init_latents = 0.18215 * init_latents
	# expand init_latents for batch_size
	init_latents = torch.cat([init_latents] * batch_size * num_images_per_prompt, dim=0)
	# get the original timestep using init_timestep
	offset = self.scheduler.config.get("steps_offset", 0)
	init_timestep = int(num_inference_steps * strength) + offset
	init_timestep = min(init_timestep, num_inference_steps)
	timesteps = self.scheduler.timesteps[-init_timestep]
	timesteps = torch.tensor([timesteps] * batch_size * num_images_per_prompt, device=self.device)
	# add noise to latents using the timesteps
	noise = torch.randn(init_latents.shape, generator=generator, device=self.device)
	init_latents = self.scheduler.add_noise(init_latents, noise, timesteps)

and

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Lines 309 to 317 in 7258dc4

	latents = init_latents
	t_start = max(num_inference_steps - init_timestep + offset, 0)
	# Some schedulers like PNDM have timesteps as arrays
	# It's more optimized to move all timesteps to correct device beforehand
	timesteps = self.scheduler.timesteps[t_start:].to(self.device)
	for i, t in enumerate(self.progress_bar(timesteps)):

Also cc @anton-l

[LowinLi]

Hi, thanks for the issue! The StableDiffusionImg2ImgPipeline only works with PyTorch models, it doesn't support onnx. For onnx you should use StableDiffusionOnnxPipeline, but note that it does not implement imag2image, but should be easy to modify to add imahe2image in it.

We will need to pass init_image as in the StableDiffusionImg2ImgPipeline and prepare the initial latents, think you pretty much copy that part into the onnx pipeline

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Lines 217 to 239 in 7258dc4

	if isinstance(init_image, PIL.Image.Image):
	init_image = preprocess(init_image)
	# encode the init image into latents and scale the latents
	init_latent_dist = self.vae.encode(init_image.to(self.device)).latent_dist
	init_latents = init_latent_dist.sample(generator=generator)
	init_latents = 0.18215 * init_latents
	# expand init_latents for batch_size
	init_latents = torch.cat([init_latents] * batch_size * num_images_per_prompt, dim=0)
	# get the original timestep using init_timestep
	offset = self.scheduler.config.get("steps_offset", 0)
	init_timestep = int(num_inference_steps * strength) + offset
	init_timestep = min(init_timestep, num_inference_steps)
	timesteps = self.scheduler.timesteps[-init_timestep]
	timesteps = torch.tensor([timesteps] * batch_size * num_images_per_prompt, device=self.device)
	# add noise to latents using the timesteps
	noise = torch.randn(init_latents.shape, generator=generator, device=self.device)
	init_latents = self.scheduler.add_noise(init_latents, noise, timesteps)

and

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Lines 309 to 317 in 7258dc4

	latents = init_latents
	t_start = max(num_inference_steps - init_timestep + offset, 0)
	# Some schedulers like PNDM have timesteps as arrays
	# It's more optimized to move all timesteps to correct device beforehand
	timesteps = self.scheduler.timesteps[t_start:].to(self.device)
	for i, t in enumerate(self.progress_bar(timesteps)):

Also cc @anton-l

Thanks, I will try your method to inference other onnx pipeline.