Running Stable Diffusion in FastAPI Does Not Release GPU Memory

[alexisrolland]

Describe the bug

I am running Stable Diffusionas as a web service using FastAPI. It runs fine, but after doing multiple inference calls, I noticed the memory of the GPU becomes full and the inference fails. It is as if the memory is not released right after doing the inference.

Reproduction

requirements.txt

--extra-index-url https://download.pytorch.org/whl/cu116
diffusers==0.3.0
fastapi==0.80.0
pydantic==1.9.2
torch==1.12.1
transformers==4.21.2
uvicorn==0.18.3

main.py

import logging
import os
import random
import time
import torch
from diffusers import StableDiffusionPipeline
from fastapi import FastAPI, HTTPException, Request
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional


# Load default logging configuration
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)

# Load Stable Diffusion model
log.info('Load Stable Diffusion model')
model_path = './models/stable-diffusion-v1-4'
pipe = StableDiffusionPipeline.from_pretrained(
    model_path,
    revision='fp16',
    torch_dtype=torch.float16
)

# Move pipeline to GPU for faster inference
pipe = pipe.to('cuda')
pipe.enable_attention_slicing()

# Declare inputs and outputs data types for the API endpoint
class Payload(BaseModel):
    prompt: str                 # String of text used to generate the images
    num_images = 1              # Number of images to be generated
    height = 512                # Height of the images to be generated
    width = 512                 # Width of the images to be generated
    seed: Optional[int] = None  # Random integer used as a seed to guide the image generator
    num_steps = 40              # Number of inference steps, results are better the more steps you use, at a cost of slower inference
    guidance_scale = 8.5        # Forces generation to better match the prompt, 7 or 8.5 give good results, results are better the larger the number is, but will be less diverse

class Response(BaseModel):
    images: List[str]
    nsfw_content_detected: List[bool]
    prompt: str
    num_images: int
    height: int
    width: int
    seed: int
    num_steps: int
    guidance_scale: float

# Create FastAPI app
log.info('Start API')
app = FastAPI(title='Stable Diffusion')
app.mount("/static", StaticFiles(directory="./static"), name="static") # Mount folder to expose generated images

# Declare imagine endpoint for inference
@app.post('/imagine', response_model=Response, description='Runs inferences with Stable Diffusion.')
def imagine(payload: Payload, request: Request):
    """The imagine function generates the /imagine endpoint and runs inferences"""

    try:
        # Check payload
        log.info(f'Payload: {payload}')

        # Default seed with a random integer if it is not provided by user
        if payload.seed is None:
            payload.seed = random.randint(-999999999, 999999999)
        generator = torch.Generator('cuda').manual_seed(payload.seed)

        # Create multiple prompts according to the number of images
        prompt = [payload.prompt] * payload.num_images

        # Run inference on GPU
        log.info('Run inference')
        with torch.autocast('cuda'):
            result = pipe(
                prompt=prompt,
                height=payload.height,
                width=payload.width,
                num_inference_steps=payload.num_steps,
                guidance_scale=payload.guidance_scale,
                generator=generator
            )
        log.info('Inference completed')

        # Save images
        images_urls = []
        for image in result.images:
            image_name = str(time.time()).replace('.', '') + '.png'
            image_path = os.path.join('static', image_name)
            image.save(image_path)
            image_url = request.url_for('static', path=image_name)
            images_urls.append(image_url)

        # Build response object
        response = {}
        response['images'] = images_urls
        response['nsfw_content_detected'] = result['nsfw_content_detected']
        response['prompt'] = payload.prompt
        response['num_images'] = payload.num_images
        response['height'] = payload.height
        response['width'] = payload.width
        response['seed'] = payload.seed
        response['num_steps'] = payload.num_steps
        response['guidance_scale'] = payload.guidance_scale

        return response

    except Exception as e:
        log.error(repr(e))
        raise HTTPException(status_code=500, detail=repr(e))

Command to run the FastAPI: uvicorn main:app --host 0.0.0.0 --port 5000

Logs

No response

System Info

• diffusers version: 0.3.0
• Platform: Linux-5.4.72-microsoft-standard-WSL2-x86_64-with-glibc2.29
• Python version: 3.8.10
• PyTorch version (GPU?): 1.12.1+cu116 (True)
• Huggingface_hub version: 0.10.1
• Transformers version: 4.21.2
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

[alexisrolland]

Screenshot showing the memory stacking-up:

[alexisrolland]

I was able to work around this issue by adding this code snippet after doing the inference... I think this should really be added to the different examples in the documentation. Kudos go to my colleague who found this from the Stable Diffusion WebUI repository.

if torch.cuda.is_available():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

[patrickvonplaten]

Thanks a lot for the issue! I'll look into this

[patrickvonplaten]

Hmmm, actually I have problems reproducing this error locally.
This could be one source of accidentally commulating GPU memory: #882 , but it shouldn't be the reason here IMO.

I don't really know what the issue could be here as we're not really accumulating tensors anywhere.

Also, we done some improvements to running stable diffusion on GPU when releasing diffusers==0.4.1 . @alexisrolland could you try to use the following code (and diffusers==0.4.1 instead) and seed whether the problem persists:

-        with torch.autocast('cuda'):
-            result = pipe(
-                prompt=prompt,
-                height=payload.height,
-                width=payload.width,
-                num_inference_steps=payload.num_steps,
-                guidance_scale=payload.guidance_scale,
-                generator=generator
-            )
+         result = pipe(
+             prompt=prompt,
+             height=payload.height,
+             width=payload.width,
+             num_inference_steps=payload.num_steps,
+             guidance_scale=payload.guidance_scale,
+             generator=generator
+         )

i.e. autocast doesn't help and diffusers=0.4.2 is fully compatible with fp16: #371

It'd be great if you could try updating to diffusers=0.4.2 to see if the problem persists.

[alexisrolland]

Hello @patrickvonplaten and thank you for answering so swiftly!
I have followed your suggestion:

• Upgraded the diffusers library, I actually upgraded straight to the latest release diffusers=0.5.1
• Removed with torch.autocast('cuda'):

It seems it did not change anything regarding the GPU memory usage. While running a bunch of requests one after another against the FastAPI, I can see the GPU memory usage growing progressively. Similarly to the screenshot I shared above.

I think this may be because each inference is ran in the same process as the FastAPI app and there is some automated caching happening at the GPU level. For comparison purpose it seems each inference on Stable Diffusion WebUI repo is ran in a separate subprocess and maybe that's why there is not such an issue there.

[patrickvonplaten]

Hmmm,

Ok I'm wondering to what extent this might just be expected when using PyTorch.

@ydshieh didn't we have a similar problem in our pytest test suite?

Also cc @anton-l @NouamaneTazi in case you have a hunch of what could be going on here

[ydshieh]

There is something similar

huggingface/accelerate#614 (comment)
(I tried locally, that issue author using Flask)

As this issue author already finds a solution using torch.cuda.empty_cache() similar to what I have done above, and also an extra
torch.cuda.ipc_collect() (I should look this, thanks), I think everything is fine.

Here are some comments:

• (It's always a good idea to try a local version to see if the issue occurs without via a web framework)
• On our side, it's not very reasonable to add torch.cuda.empty_cache() in the call method(s). We can do it at a few places (like during loading a model), but not in a call method.
• Also, the cache is managed by torch, and when you see it increases, it doesn't always mean there will be an issue. As long as that python process where torch is resides, torch will do something when new GPU memory needs to be allocated. The problem occurs when other processes want to use GPU, or the same process but other frameworks (like TensorFlow) needs GPU.

[ydshieh]

In our pytest, something similar occurs, but it is the TF tests (running after the PT tests) failed due to GPU OOM, as the PT tests didn't empty the cache.

[ydshieh]

For comparison purpose it seems each inference on Stable Diffusion WebUI repo is ran in a separate subprocess and maybe that's why there is not such an issue there.

I am just guessing: even if they use different processes, it's unlikely each inference will launch a new process. There might be multiple processes running. Potentially, after some period of time, they will replace those processes by new launched ones.

[patrickvonplaten]

Thanks a lot for the very in-detail answers @ydshieh !

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.