티스토리 뷰

728x90

 

Controlnet - Human Pose Code

ControlNet은 조건을 추가하여 Diffusion model을 제어하는 신경망 구조입니다. zero convolution이 핵심 개념이라고는 하는데, 논문에 대한 내용은 ‘Adding Conditional Control to Text-to-Image Diffusion Models’ 해당 논문을 참고하시면 좋을듯 합니다.

 

Huggingface에서 Stable Diffision과 함께 ControlNet을 결합한 모델을 사용하도록 했습니다.

 

모델 타입: Diffusion-based text-to-image generation model

 

ControlNet은 end-to-end 방식으로 작업별 조건을 학습하며 훈련 데이터셋이 작은 경우에도 robust하게 학습됩니다. 또한, ControlNet의 training은 diffusion model을 fine-tuning하는 것만큼 빠르고 개인 장치에서도 교육할 수 있다고 합니다. (colab에서도 돌아감) stable diffusion처럼 edge maps, segmentation maps, keypoints 등과 같은 조건부 입력을 가능하게 했습니다.

 

코드를 살펴보기 전, 학습되는 전체적인 구조는 이렇습니다.

 

 

입력 이미지가 들어가면, pose detector가 keypoint를 감지해서 주석처리 합니다. 그리고 prompt의 입력과 함께 저렇게 이미지를 생성해주는 것입니다.

 

1. https://github.com/patrickvonplaten/controlnet_aux 설치

pip install controlnet_aux

2. diffusers 관련 패키지를 설치:

pip install diffusers transformers accelerate

3. 실행 코드:

from PIL import Image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import torch
from controlnet_aux import OpenposeDetector
from diffusers.utils import load_image

openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')

image = load_image("<https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png>")

image = openpose(image)

controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

# Remove if you do not have xformers installed
# see <https://huggingface.co/docs/diffusers/v0.13.0/en/optimization/xformers#installing-xformers>
# for installation instructions
pipe.enable_xformers_memory_efficient_attention()

pipe.enable_model_cpu_offload()

image = pipe("chef in the kitchen", image, num_inference_steps=20).images[0]

image.save('images/chef_pose_out.png')

 

이 코드를 살펴보면,

OpenposeDetector로 사전 학습된 openpose 모델을 로드합니다. 이는 사람의 포즈를 감지하는 역할을 합니다. 그리고 이미지를 로드하고, 이 모델에 이미지를 입력하여 사람의 포즈를 추출합니다. 그 결과가 오른쪽 이미지입니다.

 

 

그 다음으로 미리 학습된 ControlNet 모델을 로드합니다. 모델이 사용하는 데이터 타입은 torch.float16 임을 알 수 있습니다.

그리고 Diffusion과 ControlNet을 결합한 파이프라인 객체 pipe를 생성합니다. controlnet 인자엔 앞에서 로드했던 ControlNet 모델을 전달합니다.

⇒ pipe( {prompt}, {이미지}, {학습 수})

  • prompt: 이미지 생성에 대한 프롬프트
  • image: 사람의 포즈 맵
  • num_inference_steps: 이미지 생성에 사용할 인퍼런스 스텝 수

 

 UniPCMultistepScheduler.from_config(pipe.scheduler.config)를 호출하여 스케줄러를 설정합니다. 스케줄러는 디프큐저 모델의 생성 과정을 조절하는 역할을 합니다.

  pipe.enable_xformers_memory_efficient_attention()를 호출하여 xformers 메모리 효율적인 어텐션을 활성화합니다. 이 기능을 사용하면 디프큐저 모델의 어텐션 계산이 최적화되어 메모리 사용량이 줄어들고 실행 속도가 향상될 수 있습니다.

 pipe.enable_model_cpu_offload()를 호출하여 모델을 CPU로 오프로드(offload)합니다. 이렇게 하면 모델의 일부 연산이 CPU에서 실행되어 GPU 메모리 사용량을 줄일 수 있습니다.

 

 

이미지가 생성됩니다. 생성된 이미지를 저장합니다.

 

ControlNet을 사용하여 openpose로 추출된 포즈가 있는 이미지를 diffusion model을 통해 조작하여 새로운 이미지를 생성하고 저장하는 과정을 보여줍니다.

 

 

출처: https://huggingface.co/lllyasviel/sd-controlnet-openpose

 

728x90
댓글