在FastAPI容器中运行Stable Diffusion时GPU显存无法释放的问题

在Docker容器里用FastAPI部署Stable Diffusion，单次推理可以正常执行，但多次调用推理接口后，GPU显存会被完全占满，导致后续推理失败，推测显存没有在推理结束后及时释放。请问有什么方法可以强制释放显存？

相关代码（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))

1. 显式清理显存缓存 #

推理完成后，手动调用PyTorch的缓存清理方法，强制释放未被使用的显存。在imagine函数的推理完成后添加以下代码：

# 清理GPU显存缓存
torch.cuda.empty_cache()
# 删除不再需要的变量，辅助垃圾回收
del result
import gc
gc.collect()

建议把这段代码放在log.info('Inference completed')之后、返回响应之前的位置。

2. 启用Diffusers的高级显存优化 #

除了已使用的enable_attention_slicing()，还可以尝试更激进的显存优化策略：

# 启用最大化注意力切片，进一步降低显存占用
pipe.enable_attention_slicing("max")
# 启用模型CPU卸载（需diffusers版本>=0.10），将模型层按需加载到GPU，用完即放回CPU
pipe.enable_model_cpu_offload()

enable_model_cpu_offload()会显著降低显存占用，但会略微增加推理时间，适合显存紧张的场景。

3. 限制并发推理请求数 #

FastAPI默认多线程处理请求，多个推理任务同时运行会快速耗尽显存。可以通过两种方式限制并发：

• 启动服务时限制工作线程和并发数：

uvicorn main:app --workers 1 --limit-concurrency 1

• 在代码中添加全局锁，确保同一时间只有一个推理任务执行：

from fastapi import BackgroundTasks
import asyncio

# 全局推理锁
inference_lock = asyncio.Lock()

@app.post('/imagine', response_model=Response, description='Runs inferences with Stable Diffusion.')
async def imagine(payload: Payload, request: Request):
    async with inference_lock:
        # 将同步推理逻辑包装为异步执行
        result = await asyncio.to_thread(run_inference, payload)
        # 后续图片保存、响应构建逻辑...

4. 定位内存泄漏点 #

可以在关键位置打印显存使用详情，排查是否有张量未被正确回收：

# 打印GPU显存使用摘要
log.info(torch.cuda.memory_summary())

重点检查result对象、生成器generator是否在使用后被及时清理，避免长期占用显存。

内容的提问来源于stack exchange，提问作者Alexis.Rolland