Concurrent Thread Failure in Image Inference: StableDiffusionPipeline

[Fantast616]

Describe the bug

When I attempt to concurrently initiate two threads for image inference in one pipe, the process fails and the following information is displayed as logs during the inference. I intend to submit a pull request for thread-safety in one inference pipe if anyone else need.

Reproduction

from diffusers import StableDiffusionPipeline
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def infer_one():
    pipe("a photo of an astronaut riding a horse on mars").images[0] 

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

Logs

Exception in thread Thread-1 (infer_one):                                                                                        | 0/50 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/workdata/A40/main.py", line 10, in infer_one
    pipe("a photo of an astronaut riding a horse on mars").images[0] 
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 710, in __call__
    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 221, in step
    return self.step_plms(model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 337, in step_plms
    prev_sample = self._get_prev_sample(sample, timestep, prev_timestep, model_output)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 371, in _get_prev_sample
    alpha_prod_t = self.alphas_cumprod[timestep]
                   ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
IndexError: index 1001 is out of bounds for dimension 0 with size 1000
  0%|                                                                                                                            | 0/50 [00:04<?, ?it/s]
Exception in thread Thread-2 (infer_one):
Traceback (most recent call last):
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/workdata/A40/main.py", line 10, in infer_one
    pipe("a photo of an astronaut riding a horse on mars").images[0] 
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/diffusers/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 710, in __call__
    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 221, in step
    return self.step_plms(model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 337, in step_plms
    prev_sample = self._get_prev_sample(sample, timestep, prev_timestep, model_output)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/workdata/yehang/projects/diffusers/src/diffusers/schedulers/scheduling_pndm.py", line 396, in _get_prev_sample
    sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff
    ~~~~~~~~~~~~~^~~~~~~~
TypeError: unsupported operand type(s) for *: 'Tensor' and 'NoneType'

System Info

• diffusers version: 0.17.0.dev0
• Platform: Linux-5.15.0-72-generic-x86_64-with-glibc2.31
• Python version: 3.11.3
• PyTorch version (GPU?): 2.0.1+cu117 (True)
• Huggingface_hub version: 0.14.1
• Transformers version: 4.29.1
• Accelerate version: 0.19.0
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patrickvonplaten]

Hey @Fantast416,

I have not yet tested SD with threads. Can you verify that the inference call works without threading to begin with?

[Fantast616]

Hey @Fantast416,

I have not yet tested SD with threads. Can you verify that the inference call works without threading to begin with?

Yes, I have verified it. if you remove the following code in my Reproduction,which means there will be only one thread run in one time, and the code will run correctly:

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

So it seems like that SD can't run two or more threads in the same time.

[Fantast616]

Hey @patrickvonplaten,
I found that the internal implementation of the diffusers framework does not consider multi-threaded inference in one pipeline object. The reason for the mentioned error is primarily due to simultaneous modifications of scheduler object's member variables by multiple threads during sampling within the schedule. I envision that if each pipeline could use an independent scheduler during inference in different threads, it should resolve this problem. Taking the StableDiffusionPipeline as an example, I propose a workaround here #3685

I would like to inquire if you have any better solution or if you believe such changes are necessary. I believe they are necessary for the deployment of open-source production services

[patrickvonplaten]

Hey @Fantast416,

That's very interesting! We indeed don't assume schedulers to work in multi-threaded environments. I think the DDIMScheduler should work for multi-threading though. Could you try whether the following works:

from diffusers import StableDiffusionPipeline, DDIMScheduler
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

def infer_one():
pipe("a photo of an astronaut riding a horse on mars").images[0]

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

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

[liyi0206]

DDIMScheduler

DDIMScheduler seems to work!

[Fantast616]

Hey,@liyi0206 @patrickvonplaten
DDIMScheduler does work, but please be careful to use it in multi-thread environments!

I have carefully examine the source code, and find that there are still thread-unsafe parts in DDIM. This is caused by the invocation of the scheduler in the 'call' function of StableDiffusionPipeline object. In the SD pipeline, 'set_timesteps' function is called first, followed by retrieving the timestep value. If there is a multi-threading conflict where timesteps are set and retrieved simultaneously, it can cause the program to deviate from the expected behavior though it may not cause exception. ( Specifically, if the following execution sequence occurs: Thread 1 calls 'set_timesteps', Thread 2 calls 'set_timesteps', Thread 1 retrieves timesteps, and Thread 2 retrieves timesteps, then Thread 1 will obtain the timestep set by Thread 2.)

Currently, I am carefully reading the source code of different schedulers and the specific invocations within StableDiffusionPipeline to figure out how to address the thread-unsafe issue internally in the sampler. I will provide a more detailed plan later.

[Fantast616]

Hello, @patrickvonplaten

• I have carefully read the code of StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, and the code of three classic samplers (PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler), and made the following research report.

• Within these two classic Pipelines, there are thread-unsafe calls that involve changing the member variables within self.scheduler, including the following functions and value retrievals:

  • .set_timesteps & .timesteps

  # 4. Prepare timesteps
  self.scheduler.set_timesteps(num_inference_steps, device=device)
  timesteps = self.scheduler.timesteps

  • .step

  with self.progress_bar(total=num_inference_steps) as progress_bar:
      for i, t in enumerate(timesteps):
          # ...
          latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
          # ...

• Therefore, we need to pay attention to whether there are any thread-unsafe parts related to these three functions in different schedulers

  • PNDMScheduler
    • .step function=> self.counter => self.counter will change in .step
    • self.counter will be set to 0 in .set_timesteps
    • multi-thread visit will let self.counter be added multiple times
  • DDIMScheduler
    • .set_timesteps and .timesteps only be called once in pipeline，but multi-thread visit will cause self.timesteps unsafe. Specifically, imagine following timeline: it will
      • ————————————————————————————————> Timeline
      • Thread-1 .set_timesteps                     .timesteps
      • Thread-2                     .set_timesteps                 .timesteps
      • Thread-1 will read the variable timesteps set by Thread-2
    • .step function will be called every iteration, but in DDIMScheduler, it will not share the counter variable, so it is multi-thread safe.
  • LMSDiscreteScheduler
    • similar to DDIMScheduler

• Based on the above idea, if we want to solve this problem within the Scheduler as you said in PR, we need to avoid distinguishing all calculations that depend on the state of a single thread member variable within the Scheduler according to different threads. This may modify a lot of code. What do you think about it?

[jasstionzyf]

@patrickvonplaten @Fantast416
I am also experience another problems caused by multi thread infer.
When i call the StableDiffusionPipeline even no problems for the Scheduler, but the image generated is in disorder.
If not use multi thread, the image generated all normal
bad image

normal image

[patrickvonplaten]

I think the easiest fix for now would be to just manually copy the scheduler in the infer_one() function, e.g. the following code snippet should work:

from diffusers import StableDiffusionPipeline
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def infer_one():
    # copy the scheduler for each thread to make it thread-safe
    pipe.scheduler = pipe.scheduler.from_config(pipe.scheduler.config)
    pipe("a photo of an astronaut riding a horse on mars").images[0] 

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

[Fantast616]

I think the easiest fix for now would be to just manually copy the scheduler in the infer_one() function, e.g. the following code snippet should work:

from diffusers import StableDiffusionPipeline
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def infer_one():
    # copy the scheduler for each thread to make it thread-safe
    pipe.scheduler = pipe.scheduler.from_config(pipe.scheduler.config)
    pipe("a photo of an astronaut riding a horse on mars").images[0] 

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

Thanks, that will work. So I think this issue should be closed.

[haoqiangyu]

I think the easiest fix for now would be to just manually copy the scheduler in the infer_one() function, e.g. the following code snippet should work:

from diffusers import StableDiffusionPipeline
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def infer_one():
    # copy the scheduler for each thread to make it thread-safe
    pipe.scheduler = pipe.scheduler.from_config(pipe.scheduler.config)
    pipe("a photo of an astronaut riding a horse on mars").images[0] 

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

I used this piece of code of yours, and while no error was thrown, I got a situation where the inference progress was frozen in those two threads.

[haoqiangyu]

I think the easiest fix for now would be to just manually copy the scheduler in the infer_one() function, e.g. the following code snippet should work:

from diffusers import StableDiffusionPipeline
import torch
import threading

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def infer_one():
    # copy the scheduler for each thread to make it thread-safe
    pipe.scheduler = pipe.scheduler.from_config(pipe.scheduler.config)
    pipe("a photo of an astronaut riding a horse on mars").images[0] 

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

I used this piece of code of yours, and while no error was thrown, I got a situation where the inference progress was frozen in those two threads.

If you use pytorch2.0.0 and above, it will cause the thread to freeze.