서버 생성

Diffusers의 파이프라인은 서버의 추론 엔진으로 사용할 수 있습니다. 여러 사용자가 동시에 요청할 수 있는 이미지 생성을 위한 동시 및 멀티스레드 요청을 지원합니다.

이 가이드에서는 서버에서 StableDiffusion3Pipeline을 사용하는 방법을 보여 드리며, 원하는 파이프라인을 자유롭게 사용할 수 있습니다.

examples/server 폴더로 이동하여 모든 의존성을 설치하는 것으로 시작합니다.

pip install .
pip install -f requirements.txt

다음 명령으로 서버를 시작합니다.

python server.py

서버는 http://localhost:8000에서 접근할 수 있습니다. 다음 명령을 사용하여 이 모델에 curl을 보낼 수 있습니다.

curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations

일부 의존성을 업그레이드해야 하는 경우 pip-tools 또는 uv를 사용할 수 있습니다. 예를 들어, 다음 명령을 사용하여 uv로 의존성을 업그레이드합니다.

uv pip compile requirements.in -o requirements.txt

서버는 FastAPI로 구축되었습니다. v1/images/generations의 엔드포인트는 다음과 같습니다.

@app.post("/v1/images/generations")
async def generate_image(image_input: TextToImageInput):
    try:
        loop = asyncio.get_event_loop()
        scheduler = shared_pipeline.pipeline.scheduler.from_config(shared_pipeline.pipeline.scheduler.config)
        pipeline = StableDiffusion3Pipeline.from_pipe(shared_pipeline.pipeline, scheduler=scheduler)
        generator = torch.Generator(device="cuda")  # or "mps", "xpu", "cpu"
        generator.manual_seed(random.randint(0, 10000000))
        output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))
        logger.info(f"output: {output}")
        image_url = save_image(output.images[0])
        return {"data": [{"url": image_url}]}
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        elif hasattr(e, 'message'):
            raise HTTPException(status_code=500, detail=e.message + traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e) + traceback.format_exc())

generate_image 함수는 async 키워드로 비동기(asynchronous)로 정의되어 있어, FastAPI가 이 함수 내에서 일어나는 작업이 즉시 결과를 반환하지 않을 수 있음을 알 수 있습니다. 함수 내에서 다른 Task를 대기(await)해야 하는 지점에 도달하면, 메인 스레드는 다른 HTTP 요청에 응답하는 작업으로 돌아갑니다. 아래 코드에서 await 키워드로 이를 확인할 수 있습니다.

output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))

이 시점에서 pipeline 함수의 실행은 새로운 스레드로 이동하며, 메인 스레드는 pipeline에서 결과가 반환될 때까지 다른 작업을 수행합니다.

이 구현의 또 다른 중요한 측면은 shared_pipeline에서 pipeline을 생성하는 것입니다. 이 작업의 목적은 GPU에 기반 모델을 두 번 이상 로드하는 것을 방지하면서, 별도의 스레드에서 실행되는 각 새 요청이 자체적인 생성기와 스케줄러를 가질 수 있도록 하는 것입니다. 특히 스케줄러는 스레드 안전성이 없으며, 여러 스레드에서 동일한 스케줄러를 사용하려고 하면 IndexError: index 21 is out of bounds for dimension 0 with size 21와 같은 오류가 발생합니다.

Update on GitHub