M2 맥북 Pro에서 Flux Schnell 로컬 모델로 고품질 이미지 생성하기

이 글에서는 Apple M2 칩 기반의 맥북 Pro 환경에서 Flux Schnell 모델을 활용하여 고품질 AI 이미지를 생성하는 방법을 단계별로 안내합니다. 특히 M2 칩의 장점을 최대한 활용하기 위한 최적화 설정과 코드에 대한 자세한 설명을 제공합니다.

1. Conda 환경 설정

먼저 Flux Schnell 모델 실행에 필요한 Python 환경을 구축하기 위해 새로운 Conda 환경을 생성하고 활성화합니다.

conda create -n flux python=3.11
conda activate flux

2. 필수 패키지 설치

다음 명령어를 사용하여 Flux Schnell 모델 실행에 필요한 필수 패키지를 설치합니다.

pip install torch==2.3.1
pip install git+https://github.com/huggingface/[email protected]
pip install transformers==4.43.3 sentencepiece==0.2.0 accelerate==0.33.0 protobuf==5.27.3

각 패키지의 역할은 다음과 같습니다.

• torch==2.3.1: PyTorch 프레임워크 (Apple Silicon 칩 호환 버전)
• git+https://github.com/huggingface/[email protected]: Diffusers 라이브러리 (특정 버전 명시)
• transformers==4.43.3: 사전 훈련된 모델 및 토크나이저 제공
• sentencepiece==0.2.0: 텍스트를 하위 단어 단위로 분할
• accelerate==0.33.0: 훈련 및 추론 가속화
• protobuf==5.27.3: 데이터 직렬화 및 통신 지원

3. Flux Schnell 모델 로딩 및 설정

다음 Python 코드는 Flux Schnell 모델을 로딩하고 M2 칩의 GPU 가속을 위한 설정을 적용합니다.

import torch
from diffusers import FluxPipeline
import diffusers


_flux_rope = diffusers.models.transformers.transformer_flux.rope


def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
    assert dim % 2 == 0, "The dimension must be even."
    if pos.device.type == "mps":
        return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device)
    else:
        return _flux_rope(pos, dim, theta)


diffusers.models.transformers.transformer_flux.rope = new_flux_rope

# Load the Flux Schnell model
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", revision="refs/pr/1", torch_dtype=torch.bfloat16
).to("mps")

# Define the prompt
prompt = "A cat holding a sign that says hello world"

# Generate the image
out = pipe(
    prompt=prompt,
    guidance_scale=0.0,
    height=1024,
    width=1024,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]

# Save the generated image
out.save("flux_image.png")

• _flux_rope 함수 재정의: M2 칩에서 MPS 백엔드를 사용하기 위해 _flux_rope 함수를 재정의합니다. 이 함수는 입력 텐서의 장치를 확인하여 CPU 또는 MPS 장치에서 적절하게 실행되도록 합니다.
• torch.bfloat16: M2 칩의 성능 향상을 위해 torch.bfloat16 데이터 유형을 사용합니다.
• .to("mps"): 모델을 Apple Silicon 칩의 GPU 가속을 위해 MPS 장치로 이동합니다.

아래는 FLUX.1-schnell 로컬 모델을 사용하여 M2 맥북 Pro에서 생성한 이미지 입니다.

_{Posted using Obsidian Steemit plugin}