Merge Stable Diffusion pipelines into one

[python273]

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

[patrickvonplaten]

Hey @python273,

According to our philosophy here: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/README.md#contribution
(see one purpose only) - I don't think it's a good idea to make all the classes into one.

However, the community might very well be interested in such a pipeline, so it would be a great addition for the community examples

[python273]

According to our philosophy here

hmm, isn't it about examples? nevermind

If you want to use a certain pipeline for a specific use case that is not yet supported, you might have to copy the pipeline file and tweak the code to your needs.

I actually combined StableDiffusionPipeline and StableDiffusionInpaintPipeline locally, but then if there are some bugfixes in the library, I would need to update my copy too.

I'm using generation + inpainting at the same time, and as I understand if I instance StableDiffusionPipeline and StableDiffusionInpaintPipeline, it will consume 2x GPU memory. So wondering if there's some way to use uploaded weights in separate pipelines

Another option might be to use a separate class for the model like:

model = StableDiffusion(...)
model = model.to("cuda")

pipe_gen = StableDiffusionPipeline(model)
pipe_inpaint = StableDiffusionInpaintPipeline(model)

[leszekhanusz]

I think that merging the pipelines into a single one in src/diffusers/pipelines/stable_diffusion is making perfect sense.

• A fix to the pipeline shoud be applied only in one place, this is the basic DRY principle.

• as @python273 said, an unified pipeline is actually needed if you want to have a single pipeline in your GPU memory for an interface for stable-diffusion doing txt to image, image to image or inpainting

@patrickvonplaten Please reconsider and allow merging the pipelines into one in the library itself. Of course in the examples folders the pipelines could and should stay simple

In the meantime, an unified pipeline version made by @nagolinc with some fixes is available here

[Inkorak]

So wondering if there's some way to use uploaded weights in separate pipelines

You can share models through pipelines. You just need to initialize one pipeline with models from another pipeline:

pipe1 = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    use_auth_token = True
).to("cuda")

pipe2 = StableDiffusionImg2ImgPipeline(vae = pipe1.vae, 
                                      text_encoder = pipe1.text_encoder,
                                      tokenizer = pipe1.tokenizer,
                                      unet = pipe1.unet,
                                      scheduler = pipe1.scheduler,
                                      safety_checker = pipe1.safety_checker,
                                      feature_extractor = pipe1.feature_extractor).to("cuda")

[python273]

oh, this works, thanks

Made a helper function to create kwargs

def extract_modules_to_dict(pipe):
    return {k: getattr(pipe, k) for k in pipe.config if not k.startswith('_')}

pipe2 = StableDiffusionInpaintPipeline(**extract_modules_to_dict(pipe))

[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

[patrickvonplaten]

Hey @Inkorak,

Feel free to add such a pipeline to our community examples https://github.com/huggingface/diffusers/tree/main/examples/community 🙂

[leszekhanusz]

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

[leszekhanusz]

Because I think it's interesting for many people I am going to maintain a unified pipeline.

to install:

pip install diffusionui

to use:

from diffusionui import StableDiffusionPipeline

then use this pipeline for either text-to-image, image-to-image or inpainting.

[patrickvonplaten]

Very cool @leszekhanusz !