FLUX, Python 및 Diffusers를 사용한 AI 기반 이미지 생성 API 서비스: 빠른 가이드

이 글에서는 Python을 사용하여 자신의 FLUX 서버를 만드는 방법을 안내해 드리겠습니다. 이 서버를 사용하면 간단한 API를 통해 텍스트 프롬프트를 기반으로 이미지를 생성할 수 있습니다. 이 서버를 개인적인 용도로 실행하든 프로덕션 애플리케이션의 일부로 배포하든 이 가이드는 시작하는 데 도움이 될 것입니다.

FLUX( Black Forest Labs 제작)는 지난 몇 달 동안 AI 이미지 생성의 세계를 강타했습니다. 많은 벤치마크에서 Stable Diffusion(이전 오픈소스의 왕)을 이겼을 뿐만 아니라 일부 지표에서 Dall-E 나 Midjourney 와 같은 독점 모델을 능가했습니다.

하지만 앱 중 하나에서 FLUX를 사용하려면 어떻게 해야 할까요? Replicate 와 같은 서버리스 호스트를 사용하는 것을 생각할 수 있지만, 이러한 호스트는 매우 빠르게 매우 비쌀 수 있으며 필요한 유연성을 제공하지 못할 수 있습니다. 여기서 사용자 정의 FLUX 서버를 만드는 것이 유용합니다.

필수 조건

코드를 살펴보기 전에 필요한 도구와 라이브러리가 설정되어 있는지 확인해 보겠습니다.

• Python: 컴퓨터에 Python 3, 바람직하게는 3.10 버전이 설치되어 있어야 합니다.
• torch : FLUX를 실행하는 데 사용할 딥러닝 프레임워크입니다.
• diffusers : FLUX 모델에 대한 액세스를 제공합니다.
• transformers : 디퓨저의 필수 종속성.
• sentencepiece : FLUX 토크나이저를 실행하는 데 필요합니다.
• protobuf : FLUX 실행에 필요함
• accelerate : 어떤 경우에는 FLUX 모델을 더 효율적으로 로드하는 데 도움이 됩니다.
• fastapi : 이미지 생성 요청을 수용할 수 있는 웹 서버를 만드는 프레임워크입니다.
• uvicorn : FastAPI 서버를 실행하는 데 필요합니다.
• psutil : 컴퓨터에 얼마나 많은 RAM이 있는지 확인할 수 있습니다.

다음 명령을 실행하여 모든 라이브러리를 설치할 수 있습니다: pip install torch diffusers transformers sentencepiece protobuf accelerate fastapi uvicorn .

M1 또는 M2 칩이 있는 Mac을 사용하는 경우 최적의 성능을 위해 PyTorch를 Metal로 설정해야 합니다. 진행하기 전에 공식 PyTorch with Metal 가이드를 따르세요.

GPU 디바이스에서 FLUX를 실행하려면 최소 12GB의 VRAM이 있어야 합니다. 또는 CPU/MPS에서 실행하려면 최소 12GB의 RAM이 있어야 합니다(이는 더 느립니다).

1단계: 환경 설정

먼저, 우리가 사용하는 하드웨어에 따라 추론을 실행할 적절한 장치를 선택하여 스크립트를 시작해 보겠습니다.

 device = 'cuda' # can also be 'cpu' or 'mps' import os # MPS support in PyTorch is not yet fully implemented if device == 'mps': os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" import torch if device == 'mps' and not torch.backends.mps.is_available(): raise Exception("Device set to MPS, but MPS is not available") elif device == 'cuda' and not torch.cuda.is_available(): raise Exception("Device set to CUDA, but CUDA is not available")

cpu , cuda (NVIDIA GPU의 경우) 또는 mps (Apple의 Metal Performance Shaders의 경우)를 지정할 수 있습니다. 그런 다음 스크립트는 선택한 장치를 사용할 수 있는지 확인하고 그렇지 않으면 예외를 발생시킵니다.

2단계: FLUX 모델 로드

다음으로, FLUX 모델을 로드합니다. fp16 정밀도로 모델을 로드하면 품질이 크게 떨어지지 않고 메모리를 약간 절약할 수 있습니다.

이 시점에서 FLUX 모델이 게이트되어 있으므로 HuggingFace로 인증하라는 메시지가 표시될 수 있습니다. 성공적으로 인증하려면 HuggingFace 계정을 만들고 모델 페이지로 이동하여 약관에 동의한 다음 계정 설정에서 HuggingFace 토큰을 만들고 HF_TOKEN 환경 변수로 컴퓨터에 추가해야 합니다.

 from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline import psutil model_name = "black-forest-labs/FLUX.1-dev" print(f"Loading {model_name} on {device}") pipeline = FluxPipeline.from_pretrained( model_name, # Diffusion models are generally trained on fp32, but fp16 # gets us 99% there in terms of quality, with just half the (V)RAM torch_dtype=torch.float16, # Ensure we don't load any dangerous binary code use_safetensors=True # We are using Euler here, but you can also use other samplers scheduler=FlowMatchEulerDiscreteScheduler() ).to(device)

여기서는, 우리는 디퓨저 라이브러리를 사용하여 FLUX 모델을 로드하고 있습니다. 우리가 사용하는 모델은 black-forest-labs/FLUX.1-dev 이며, fp16 정밀도로 로드되었습니다.

또한 더 빠른 추론을 제공하지만 덜 자세한 이미지를 출력하는 FLUX Schnell이라는 타임스텝 증류 모델과 폐쇄 소스인 FLUX Pro 모델도 있습니다. 여기서는 Euler 스케줄러를 사용하지만, 이를 실험해 볼 수 있습니다. 스케줄러에 대한 자세한 내용은 여기 에서 읽을 수 있습니다. 이미지 생성은 리소스를 많이 소모할 수 있으므로, 특히 CPU나 메모리가 제한된 장치에서 실행할 때 메모리 사용을 최적화하는 것이 중요합니다.

 # Recommended if running on MPS or CPU with < 64 GB of RAM total_memory = psutil.virtual_memory().total total_memory_gb = total_memory / (1024 ** 3) if (device == 'cpu' or device == 'mps') and total_memory_gb < 64: print("Enabling attention slicing") pipeline.enable_attention_slicing()

이 코드는 사용 가능한 총 메모리를 확인하고 시스템의 RAM이 64GB 미만인 경우 어텐션 슬라이싱을 활성화합니다. 어텐션 슬라이싱은 이미지 생성 중 메모리 사용량을 줄여 주는데, 이는 리소스가 제한된 기기에 필수적입니다.

3단계: FastAPI를 사용하여 API 생성

다음으로, 이미지 생성 API를 제공하는 FastAPI 서버를 설정하겠습니다.

 from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field, conint, confloat from fastapi.middleware.gzip import GZipMiddleware from io import BytesIO import base64 app = FastAPI() # We will be returning the image as a base64 encoded string # which we will want compressed app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

FastAPI는 Python으로 웹 API를 구축하는 데 인기 있는 프레임워크입니다. 이 경우, 이미지 생성 요청을 수락할 수 있는 서버를 만드는 데 사용합니다. 또한 GZip 미들웨어를 사용하여 응답을 압축하는데, 이는 특히 base64 형식으로 이미지를 다시 보낼 때 유용합니다.

운영 환경에서는 생성된 이미지를 S3 버킷이나 다른 클라우드 스토리지에 저장하고 base64로 인코딩된 문자열 대신 URL을 반환하여 CDN 및 기타 최적화를 활용할 수 있습니다.

4단계: 요청 모델 정의

이제 API가 수락할 요청에 대한 모델을 정의해야 합니다.

 class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch")

이 GenerateRequest 모델은 이미지를 생성하는 데 필요한 매개변수를 정의합니다. prompt 필드는 만들고자 하는 이미지의 텍스트 설명입니다. 다른 필드에는 이미지 크기, 추론 단계 수, 배치 크기가 포함됩니다.

5단계: 이미지 생성 엔드포인트 생성

이제 이미지 생성 요청을 처리할 엔드포인트를 만들어 보겠습니다.

 @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URLs instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, }

이 엔드포인트는 이미지 생성 프로세스를 처리합니다. 먼저 FLUX에서 요구하는 대로 높이와 너비가 8의 배수인지 확인합니다. 그런 다음 제공된 프롬프트를 기반으로 이미지를 생성하여 base64로 인코딩된 문자열로 반환합니다.

6단계: 서버 시작

마지막으로 스크립트가 실행될 때 서버를 시작하는 코드를 추가해 보겠습니다.

 @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)

이 코드는 FastAPI 서버를 포트 8000에서 시작하여 http://localhost:8000 에서 접근할 수 있을 뿐만 아니라 0.0.0.0 바인딩 덕분에 호스트 머신의 IP 주소를 사용하는 동일 네트워크의 다른 장치에서도 접근할 수 있게 합니다.

7단계: 로컬에서 서버 테스트

이제 FLUX 서버가 가동되었으므로 테스트할 시간입니다. HTTP 요청을 만드는 명령줄 도구인 curl 사용하여 서버와 상호 작용할 수 있습니다.

 curl -X POST "http://localhost:8000/" \ -H "Content-Type: application/json" \ -d '{ "prompt": "A futuristic cityscape at sunset", "seed": 42, "height": 1024, "width": 1024, "cfg": 3.5, "steps": 50, "batch_size": 1 }' | jq -r '.images[0]' | base64 -d > test.png

이 명령은 curl , jq 및 base64 유틸리티가 설치된 UNIX 기반 시스템에서만 작동합니다. FLUX 서버를 호스팅하는 하드웨어에 따라 완료하는 데 몇 분이 걸릴 수도 있습니다.

결론

축하합니다! Python을 사용하여 자신의 FLUX 서버를 성공적으로 만들었습니다. 이 설정을 사용하면 간단한 API를 통해 텍스트 프롬프트를 기반으로 이미지를 생성할 수 있습니다. 기본 FLUX 모델의 결과에 만족하지 못하는 경우 특정 사용 사례에서 더 나은 성능을 위해 모델을 미세 조정하는 것을 고려할 수 있습니다 .

전체 코드

이 가이드에 사용된 전체 코드는 아래에서 확인할 수 있습니다.

 device = 'cuda' # can also be 'cpu' or 'mps' import os # MPS support in PyTorch is not yet fully implemented if device == 'mps': os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" import torch if device == 'mps' and not torch.backends.mps.is_available(): raise Exception("Device set to MPS, but MPS is not available") elif device == 'cuda' and not torch.cuda.is_available(): raise Exception("Device set to CUDA, but CUDA is not available") from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline import psutil model_name = "black-forest-labs/FLUX.1-dev" print(f"Loading {model_name} on {device}") pipeline = FluxPipeline.from_pretrained( model_name, # Diffusion models are generally trained on fp32, but fp16 # gets us 99% there in terms of quality, with just half the (V)RAM torch_dtype=torch.float16, # Ensure we don't load any dangerous binary code use_safetensors=True, # We are using Euler here, but you can also use other samplers scheduler=FlowMatchEulerDiscreteScheduler() ).to(device) # Recommended if running on MPS or CPU with < 64 GB of RAM total_memory = psutil.virtual_memory().total total_memory_gb = total_memory / (1024 ** 3) if (device == 'cpu' or device == 'mps') and total_memory_gb < 64: print("Enabling attention slicing") pipeline.enable_attention_slicing() from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field, conint, confloat from fastapi.middleware.gzip import GZipMiddleware from io import BytesIO import base64 app = FastAPI() # We will be returning the image as a base64 encoded string # which we will want compressed app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7) class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch") @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URL's instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, } @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)