Merge Stable Diffusion pipelines into one

Is your feature request related to a problem? Please describe. It would be nice to be able to use text2img & img2img & inpaint with the same instanced pipeline

Describe the solution you’d like pipe.generate(...) pipe.image_to_image(...) pipe.inpaint(...)

Describe alternatives you’ve considered It’s possible to call StableDiffusionInpaintPipeline.__call__(pipe, ...), but it might break at some point if classes are not compatible

[leszekhanusz]

@Inkorak With a few changes I was able to make it work.

[Inkorak]

@python273 You can also write something like this. But I haven’t tried it myself yet.

class  StableDiffusionUnifiedPipeline(DiffusionPipeline):
    def __init__(
        self,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        unet: UNet2DConditionModel,
        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
        safety_checker: StableDiffusionSafetyChecker,
        feature_extractor: CLIPFeatureExtractor,
    ):
        super().__init__()
        text2img = StableDiffusionPipeline(vae=vae,
                                          text_encoder=text_encoder,
                                          tokenizer=tokenizer,
                                          unet=unet,
                                          scheduler=scheduler,
                                          safety_checker=safety_checker,
                                          feature_extractor=feature_extractor)
        img2img = StableDiffusionImg2ImgPipeline(vae=vae,
                                          text_encoder=text_encoder,
                                          tokenizer=tokenizer,
                                          unet=unet,
                                          scheduler=scheduler,
                                          safety_checker=safety_checker,
                                          feature_extractor=feature_extractor)
        inpaint =  StableDiffusionInpaintPipeline(vae=vae,
                                          text_encoder=text_encoder,
                                          tokenizer=tokenizer,
                                          unet=unet,
                                          scheduler=scheduler,
                                          safety_checker=safety_checker,
                                          feature_extractor=feature_extractor)
        
        self.register_modules(
            text2img=text2img,
            img2img=img2img,
            inpaint=inpaint,
        )

    def __call__(
        self,
        prompt: Union[str, List[str]],
        init_image: Union[torch.FloatTensor, PIL.Image.Image],
        mask_image: Union[torch.FloatTensor, PIL.Image.Image],
        strength: float = 0.8,
        height: Optional[int] = 512,
        width: Optional[int] = 512,
        num_inference_steps: Optional[int] = 50,
        guidance_scale: Optional[float] = 7.5,
        eta: Optional[float] = 0.0,
        generator: Optional[torch.Generator] = None,
        latents: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        **kwargs,
    ):
      if init_image is None and mask_image is None:
        result = self.text2img(prompt =prompt,
                    height = height,
                    width = width,
                    num_inference_steps = num_inference_steps,
                    guidance_scale = guidance_scale,
                    eta = eta,
                    generator = generator,
                    latents = latents,
                    output_type = output_type,
                    **kwargs,)
      elif init_image is not None and mask_image is None:
        result = self.img2img(prompt =prompt,
                    init_image = init_image,
                    strength = strength,
                    num_inference_steps = num_inference_steps,
                    guidance_scale = guidance_scale,
                    eta = eta,
                    generator = generator,
                    output_type = output_type,
                    **kwargs,)
      else:
        result = self.inpaint(prompt =prompt,
                    init_image = init_image,
                    mask_image = mask_image,
                    strength = strength,
                    num_inference_steps = num_inference_steps,
                    guidance_scale = guidance_scale,
                    eta = eta,
                    generator = generator,
                    output_type = output_type,
                    **kwargs,)
      return result