Activating attention slicing leads to black images when running diffusion more than once

I get a black image when I run a pipeline the second time and enable_attention_slicing() is enabled.

I am following those instructions:

And this is the code I am using right now in a Jupyter notebook:

from tqdm import tqdm

from diffusers import DiffusionPipeline

from diffusers import DPMSolverMultistepScheduler

import torch

device = "mps" # cuda

model_id = "runwayml/stable-diffusion-v1-5"

pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True, )

pipeline.enable_attention_slicing() #

pipeline = pipeline.to(device)

pipeline.safety_checker = None

pipeline.requires_safety_checker = False

pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)

generator = torch.Generator(device).manual_seed(0)

prompt = "portrait photo of a old warrior chief"

image = pipeline(prompt, generator=generator, num_inference_steps=5).images[0]

image

As soon as I run this a second time or e.g. when I run batch image generation I either get the NSWF warning. If I deactivate it by e.g. setting safety_checker = None I get black results only.

I have to restart the notebook to make it work again. But then again, only the first run works.

When I remove pipeline.enable_attention_slicing() everything works fine.

Is this a bug or is there something wrong with my config?

//NB
seed is fixed and steps is set to 5 to ease reproduction, changing that doesn’t help

Found a couple of possible solutions and it seems like the attentio_slicing process is a little buggy:

the main “bug” is probably:
After enable_attention_slicing() you cannot re-run the inference again. Cuda provides a reset method that probably solves that. But that does not work on MPS (Apple Silicon). I am trying to find another way to quickly reset the device, instead of fully reloading everything.

But there are also a couple of depdencies:

  • don’t pass cross_attention_kwargs or guidance_scale to the pipelin method when running the inference, even if set to None, they will result in the black image
  • when doing trial & error, not just re-run the script, re-start your Jupyter (if you use Jupyter)

Not sure if they are connected, but:

  • make sure that a compatible AutoEncoder is selected (or none) (it’s not in the code above, found out while playing around)
  • make sure that a compatbile scheduler/solver is selected (e.g. DPMMultiStep seems to be incompatible with SD15, better use DDPMScheduler)

FWIWF:
Tricks like

import gc
with torch.no_grad():
    torch.mps.empty_cache()
gc.collect()

or

del pipeline

won’t help.

You have to reload everything, like your Jupyter notebook or the Gradio app.