AutoencoderKL multi thread issues

[jasstionzyf]

Describe the bug

Multi pipelines use same AutoencoderKL run concurrency lead to wired image like black.

Reproduction

from diffusers import StableDiffusionPipeline
import torch
import threading
import copy
import numpy as np
import torch
from diffusers import (

StableDiffusionPipeline,
UNet2DConditionModel,
AutoencoderKL, UNet2DConditionModel, PNDMScheduler,

DPMSolverSinglestepScheduler
    


)
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
import os
from compel import Compel
from IPython.display import display
from PIL import Image

modelPath="/root/.cache/soujpg/models/diffusers/sd-1.5/"

vae= AutoencoderKL.from_pretrained(modelPath+"vae/",torch_dtype=torch.float16)
unet= UNet2DConditionModel.from_pretrained(modelPath, subfolder="unet",torch_dtype=torch.float16)
text_encoder= CLIPTextModel.from_pretrained(modelPath, subfolder="text_encoder",torch_dtype=torch.float16)
tokenizer= CLIPTokenizer.from_pretrained(modelPath, subfolder="tokenizer",torch_dtype=torch.float16)
scheduler=DPMSolverSinglestepScheduler.from_pretrained(modelPath, subfolder="scheduler",torch_dtype=torch.float16) 
safety_checker=None
feature_extractor=None






def infer_one():


    pipelineKwargs={
        "unet":unet,
        "text_encoder":text_encoder,
        "vae":vae,
        "tokenizer":tokenizer,
        "feature_extractor":feature_extractor,
        "safety_checker":safety_checker,
        "scheduler":copy.deepcopy(scheduler),}
    pipeline=StableDiffusionPipeline(**pipelineKwargs)
    pipeline=pipeline.to("cuda")
    rs=pipeline("a photo of an astronaut riding a horse on mars").images[0] 
    rs.save(shortuuid.uuid()+".png")
    

x1 = threading.Thread(target=infer_one) 
x1.start()

x2 = threading.Thread(target=infer_one) 
x2.start()

Logs

No response

System Info

diffusers version: 0.27.0
Platform: Linux-6.1.58+-x86_64-with-glibc2.35
Python version: 3.11.4
PyTorch version (GPU?): 2.2.1+cu122 (True)
Huggingface_hub version: 0.20.3
Transformers version: 4.38.2
Accelerate version: 0.28.0
xFormers version: not installed
Using GPU in script?: yes, RTX 4080
Using distributed or parallel set-up in script?: yes

Who can help?

@yiyixuxu @DN6 @sayakpaul

[jasstionzyf]

But if i use deep copy vae, then works normal. generated image is OK.

[jasstionzyf]

Another observation is that even i deep copy unet and vae, with three threads run concurrencym then all black images.
but with two threads run is ok.

[sayakpaul]

I don't think this an issue specific to diffusers. If you can trace the issue back to the library, of course, feel free to open it as an issue. Otherwise, I think this should be Discussions.

[jasstionzyf]

@sayakpaul
Yes, i can not sure this is related with diffusers libary.
But it seems that i can not use multi diffusers pipeline concurrency with resuable modules like vae, unet.

Also i think multi pipeline with shared vae run concurrency can improve interfence performence.

Or are there somethings i miss to use diffusers in multi threads environments?

[sayakpaul]

Pipeline concurrency is not something we have looked into. We do have support for distributed inference which you can find out from our docs.

[jasstionzyf]

I means one gpu run multi diffusers pipeline concurrency.

[sayakpaul]

Here's the code I am trying on a single A100:

from diffusers import StableDiffusionPipeline
import torch
import threading
import copy
import torch
from diffusers import (
    StableDiffusionPipeline,
    UNet2DConditionModel,
    AutoencoderKL,
    UNet2DConditionModel,
    DPMSolverSinglestepScheduler,
)
from transformers import CLIPTextModel, CLIPTokenizer
import uuid

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

vae = AutoencoderKL.from_pretrained(modelPath, subfolder="vae", torch_dtype=torch.float16)
unet = UNet2DConditionModel.from_pretrained(modelPath, subfolder="unet", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained(modelPath, subfolder="text_encoder", torch_dtype=torch.float16)
tokenizer = CLIPTokenizer.from_pretrained(modelPath, subfolder="tokenizer", torch_dtype=torch.float16)
scheduler = DPMSolverSinglestepScheduler.from_pretrained(modelPath, subfolder="scheduler", torch_dtype=torch.float16)
safety_checker = None
feature_extractor = None


def infer_one():
    pipelineKwargs = {
        "unet": unet,
        "text_encoder": text_encoder,
        "vae": vae,
        "tokenizer": tokenizer,
        "feature_extractor": feature_extractor,
        "safety_checker": safety_checker,
        "scheduler": copy.deepcopy(scheduler),
    }
    pipeline = StableDiffusionPipeline(**pipelineKwargs)
    pipeline = pipeline.to("cuda")
    rs = pipeline("a photo of an astronaut riding a horse on mars").images[0]
    rs.save(str(uuid.uuid4()) + ".png")


x1 = threading.Thread(target=infer_one)
x1.start()

x2 = threading.Thread(target=infer_one)
x2.start()

It's leading segmentation fault.

[sayakpaul]

Another inherent problem while achieving concurrency could be that a diffusion process by nature is very sequential in nature. So, too much concurrency can mess up with the intermediate states (such as the ones in the scheduler). As such, I don't think they are all fully thread-safe.

[sayakpaul]

Okay, I investigated a couple of things.

My code modified from the provided one

from diffusers import StableDiffusionPipeline
import torch
import threading
import copy
import torch
from diffusers import (
    StableDiffusionPipeline,
    UNet2DConditionModel,
    AutoencoderKL,
    UNet2DConditionModel,
    DPMSolverSinglestepScheduler,
)
from transformers import CLIPTextModel, CLIPTokenizer
import uuid

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

vae = AutoencoderKL.from_pretrained(modelPath, subfolder="vae", torch_dtype=torch.float16)
unet = UNet2DConditionModel.from_pretrained(modelPath, subfolder="unet", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained(modelPath, subfolder="text_encoder", torch_dtype=torch.float16)
tokenizer = CLIPTokenizer.from_pretrained(modelPath, subfolder="tokenizer", torch_dtype=torch.float16)
scheduler = DPMSolverSinglestepScheduler.from_pretrained(modelPath, subfolder="scheduler", torch_dtype=torch.float16)
safety_checker = None
feature_extractor = None

pipelineKwargs = {
    "unet": unet,
    "text_encoder": text_encoder,
    "vae": vae,
    "tokenizer": tokenizer,
    "feature_extractor": feature_extractor,
    "safety_checker": safety_checker,
    "scheduler": copy.deepcopy(scheduler),
}
pipeline = StableDiffusionPipeline(**pipelineKwargs)
pipeline = pipeline.to("cuda")

def infer_one():
    rs = pipeline("a photo of an astronaut riding a horse on mars").images[0]
    rs.save(str(uuid.uuid4()) + ".png")


x1 = threading.Thread(target=infer_one)
x1.start()

x2 = threading.Thread(target=infer_one)
x2.start()

This leads to:

    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
  File "/home/sayak/diffusers/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py", line 915, in step
    return func(*args, **kwargs)
  File "/home/sayak/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 993, in __call__
    order = self.order_list[self.step_index]
IndexError: list index out of range
    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
  File "/home/sayak/diffusers/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py", line 915, in step
    order = self.order_list[self.step_index]
IndexError: list index out of range

Short:

    order = self.order_list[self.step_index]
IndexError: list index out of range

TL;DR: Our schedulers are NOT thread-safe.

Long description

The problem happens because with step(), this happens in a non thread-safe way:

diffusers/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

Line 976 in e29f16c

self._step_index += 1

This is what causes the step_index to change in an unexpected manner.

To remedy this problem (which is a bit out of the scope for us the moment), we first need to ensure that the step_index increment happens in a thread-safe way ensuring that only one thread can update this value. For this kind of mutual exclusion to happen, we need to make the corresponding scheduler class accept a global lock as an argument.

And then the threading logic needs to be crafted accordingly. I wrote the following script to simulate a similar scenario.

import threading
from diffusers import DPMSolverSinglestepScheduler
import torch
from tqdm.auto import tqdm


global_lock = threading.Lock()

class SchedulerStepSimulator:
    def __init__(self, num_inference_steps=10, lock=None):
        scheduler_kwargs = {
            "beta_end": 0.012,
            "beta_schedule": "scaled_linear",
            "beta_start": 0.00085,
            "num_train_timesteps": 1000,
            "trained_betas": None,
        }
        self.scheduler = DPMSolverSinglestepScheduler(**scheduler_kwargs, lock=lock)
        self.scheduler.set_timesteps(num_inference_steps)
        self.timesteps = self.scheduler.timesteps
        self.num_inference_steps = num_inference_steps
        print(self.timesteps, self.num_inference_steps)
    
    def run(self):
        model_output = torch.randn(1, 4, 64, 64)
        sample = torch.randn(1, 4, 64, 64)
        with tqdm(total=self.num_inference_steps) as pbar:
            for i, t in enumerate(self.timesteps): 
                _ = self.scheduler.step(model_output, t, sample, return_dict=False)[0]
                if i == len(self.timesteps) - 1 or ((i + 1) % self.scheduler.order == 0):
                    pbar.update()

counter1 = SchedulerStepSimulator(lock=global_lock)
counter2 = SchedulerStepSimulator(lock=global_lock)

# Create and start threads
thread1 = threading.Thread(target=counter1.run)
thread2 = threading.Thread(target=counter2.run)

print("Starting threads")
thread1.start()
thread2.start()

# Wait for both threads to complete
thread1.join()
thread2.join()
print("Both threads finished execution")

This ensures that the step_index increment happens in a thread-safe manner.

In order for the above code to execute, you also need the following patch applied to the scheduler implementation.

diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
index 215d94a78..5d754c8f3 100644
--- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
+++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
@@ -16,6 +16,7 @@
 
 import math
 from typing import List, Optional, Tuple, Union
+from contextlib import nullcontext
 
 import numpy as np
 import torch
@@ -156,6 +157,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
         final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
         lambda_min_clipped: float = -float("inf"),
         variance_type: Optional[str] = None,
+        lock=None,
     ):
         if algorithm_type == "dpmsolver":
             deprecation_message = "algorithm_type `dpmsolver` is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
@@ -212,6 +214,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
         self._step_index = None
         self._begin_index = None
         self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self._lock = lock or nullcontext()
 
     def get_order_list(self, num_inference_steps: int) -> List[int]:
         """
@@ -908,6 +911,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
             self.model_outputs[i] = self.model_outputs[i + 1]
         self.model_outputs[-1] = model_output
 
+        print(f"self.step_index: {self.step_index}")
         order = self.order_list[self.step_index]
 
         #  For img2img denoising might start with order>1 which is not possible
@@ -922,7 +926,8 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin):
         prev_sample = self.singlestep_dpm_solver_update(self.model_outputs, sample=self.sample, order=order)
 
         # upon completion increase step index by one
-        self._step_index += 1
+        with self._lock:
+            self._step_index += 1
 
         if not return_dict:
             return (prev_sample,)

Hope that helps.

Ccing @yiyixuxu @BenjaminBossan for their comments too if they're interested :)

[jasstionzyf]

Here's the code I am trying on a single A100:

from diffusers import StableDiffusionPipeline
import torch
import threading
import copy
import torch
from diffusers import (
    StableDiffusionPipeline,
    UNet2DConditionModel,
    AutoencoderKL,
    UNet2DConditionModel,
    DPMSolverSinglestepScheduler,
)
from transformers import CLIPTextModel, CLIPTokenizer
import uuid

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

vae = AutoencoderKL.from_pretrained(modelPath, subfolder="vae", torch_dtype=torch.float16)
unet = UNet2DConditionModel.from_pretrained(modelPath, subfolder="unet", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained(modelPath, subfolder="text_encoder", torch_dtype=torch.float16)
tokenizer = CLIPTokenizer.from_pretrained(modelPath, subfolder="tokenizer", torch_dtype=torch.float16)
scheduler = DPMSolverSinglestepScheduler.from_pretrained(modelPath, subfolder="scheduler", torch_dtype=torch.float16)
safety_checker = None
feature_extractor = None


def infer_one():
    pipelineKwargs = {
        "unet": unet,
        "text_encoder": text_encoder,
        "vae": vae,
        "tokenizer": tokenizer,
        "feature_extractor": feature_extractor,
        "safety_checker": safety_checker,
        "scheduler": copy.deepcopy(scheduler),
    }
    pipeline = StableDiffusionPipeline(**pipelineKwargs)
    pipeline = pipeline.to("cuda")
    rs = pipeline("a photo of an astronaut riding a horse on mars").images[0]
    rs.save(str(uuid.uuid4()) + ".png")


x1 = threading.Thread(target=infer_one)
x1.start()

x2 = threading.Thread(target=infer_one)
x2.start()

It's leading segmentation fault.

This is actually wired for we have deep copy scheduler and for each thread, also my case at 4080 was black image.
I seems differnet environments when use multi threads run invidual pipeline may lead to different bad results

[bghira]

why not simply batch the inputs? different sample sizes?

[jasstionzyf]

why not simply batch the inputs? different sample sizes?

yes, different sample sizes , different loars and controlnets .... even with same unet and vae