Custom Diffusion

Custom Diffusion은 이미지 생성 모델을 개인화하기 위한 학습 기법입니다. Textual Inversion, DreamBooth, LoRA와 마찬가지로 Custom Diffusion은 소수의 (~4-5개) 예시 이미지만 필요합니다. 이 기법은 cross-attention 레이어의 가중치만 학습하는 방식으로 동작하며, 새로 학습된 개념을 나타내기 위해 특수한 단어를 사용합니다. Custom Diffusion은 여러 개념을 동시에 학습할 수도 있다는 점에서 독특합니다.

vRAM이 제한된 GPU에서 학습하는 경우, --enable_xformers_memory_efficient_attention을 사용하여 xFormers를 활성화하면 vRAM 요구 사항(16GB)을 낮추고 더 빠르게 학습할 수 있습니다. 메모리를 더 절약하려면 학습 인자에 --set_grads_to_none을 추가하여 그라디언트를 0 대신 None으로 설정하세요(이 옵션은 일부 문제를 일으킬 수 있으므로 문제가 발생하면 이 매개변수를 제거해 보세요).

이 가이드에서는 train_custom_diffusion.py 스크립트를 살펴보고, 이를 더 잘 이해하며 자신의 사용 사례에 맞게 수정하는 방법을 안내합니다.

…

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

학습 스크립트가 포함된 예제 폴더로 이동하고 필요한 의존성을 설치합니다:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 Accelerate는 여러 GPU/TPU 또는 혼합 정밀도로 학습하는 것을 돕는 라이브러리입니다. 하드웨어와 환경에 따라 학습 설정을 자동으로 구성합니다. 자세한 내용은 🤗 Accelerate Quick tour를 참고하세요.

🤗 Accelerate 환경을 초기화합니다:

accelerate config

어떤 구성도 선택하지 않고 기본 🤗 Accelerate 환경을 설정하려면:

accelerate config default

또는 노트북과 같이 인터랙티브 셸을 지원하지 않는 환경에서는 다음을 사용할 수 있습니다:

from accelerate.utils import write_basic_config

write_basic_config()

마지막으로, 자체 데이터셋으로 모델을 학습하려면 학습용 데이터셋 만들기 가이드를 참고하여 학습 스크립트와 호환되는 데이터셋을 만드는 방법을 배워보세요.

다음 섹션에서는 스크립트를 수정하는 방법을 이해하는 데 중요한 학습 스크립트의 일부 부분을 강조합니다. 그러나 스크립트의 모든 측면을 자세히 다루지는 않습니다. 더 자세히 알고 싶으시다면 스크립트를 읽어보시고, 질문이나 우려 사항이 있으시면 알려주시기 바랍니다.

스크립트 매개변수

학습 스크립트에는 학습 실행을 사용자 정의하는 데 도움이 되는 모든 매개변수가 포함되어 있습니다. 이들은 parse_args() 함수에서 찾을 수 있습니다. 이 함수에는 기본값이 제공되지만, 원하시는 경우 학습 명령에서 자체 값을 설정할 수도 있습니다.

예를 들어, 입력 이미지의 해상도를 변경하려면:

accelerate launch train_custom_diffusion.py \
  --resolution=256

기본 매개변수 중 상당수는 DreamBooth 학습 가이드에 설명되어 있으므로, 이 가이드에서는 Custom Diffusion에 고유한 매개변수에 초점을 맞춥니다:

--freeze_model: cross-attention 레이어의 key와 value 파라미터를 동결합니다. 기본값은 crossattn_kv이며, cross-attention 레이어의 모든 파라미터를 학습하려면 crossattn으로 설정할 수 있습니다.
--concepts_list: 여러 개념을 학습하려면 개념을 포함하는 JSON 파일의 경로를 제공합니다
--modifier_token: 학습된 개념을 나타내는 데 사용되는 특수 단어
--initializer_token: modifier_token의 임베딩을 초기화하는 데 사용되는 특수 단어

사전 보존 손실

Prior preservation loss는 모델이 스스로 생성한 샘플을 활용하여 더 다양한 이미지를 생성하는 방법을 학습하도록 돕는 기법입니다. 이러한 생성된 샘플 이미지는 사용자가 제공한 이미지와 동일한 클래스에 속하므로, 모델이 해당 클래스에 대해 학습한 내용을 유지하고 이미 알고 있는 클래스 지식을 활용하여 새로운 구성을 만들어내는 데 도움이 됩니다.

prior preservation loss의 많은 매개변수는 DreamBooth 학습 가이드에 설명되어 있습니다.

정규화

Custom Diffusion은 과적합을 방지하기 위해 소수의 실제 이미지로 대상 이미지를 학습하는 것을 포함합니다. 단 몇 장의 이미지로만 학습하는 경우이므로, 이는 쉽게 수행할 수 있습니다. clip_retrieval을 사용하여 실제 이미지 200장을 다운로드하세요. class_prompt는 대상 이미지와 동일한 카테고리여야 합니다. 이 이미지들은 class_data_dir에 저장됩니다.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

정규화를 활성화하려면 다음 매개변수를 추가하세요:

--with_prior_preservation: prior preservation loss를 사용할지 여부
--prior_loss_weight: 모델에 대한 사전 보존 손실의 영향을 제어합니다
--real_prior: 과적합을 방지하기 위해 실제 이미지의 작은 집합을 사용할지 여부

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

학습 스크립트

Custom Diffusion 학습 스크립트의 많은 코드는 DreamBooth 스크립트와 유사합니다. 이 가이드는 대신 Custom Diffusion와 관련된 코드에 초점을 맞춥니다.

Custom Diffusion 학습 스크립트에는 두 가지 데이터셋 클래스가 있습니다:

CustomDiffusionDataset: 학습을 위해 이미지, 클래스 이미지, 프롬프트를 전처리합니다
PromptDataset: 클래스 이미지 생성을 위한 프롬프트를 준비합니다

다음으로, modifier_token은 토크나이저에 추가되고 토큰 ID로 변환되며, 새로운 modifier_token을 반영하도록 토큰 임베딩의 크기가 조정됩니다. 그런 다음 modifier_token의 임베딩은 initializer_token의 임베딩으로 초기화됩니다. 텍스트 인코더의 모든 매개변수는 동결되지만, 모델이 개념과 연관 지어 학습하려는 대상이 토큰 임베딩이므로 토큰 임베딩만 예외입니다.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

이제 Custom Diffusion weights를 attention 레이어에 추가해야 합니다. 이는 attention weights의 shape와 size를 올바르게 설정하고 각 UNet 블록에 적절한 수의 attention processors를 설정하는 데 매우 중요한 단계입니다.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

optimizer는 크로스 어텐션 레이어의 매개변수를 업데이트하도록 초기화됩니다:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

training loop에서는 학습하려는 개념에 해당하는 임베딩만 업데이트하는 것이 중요합니다. 즉, 나머지 모든 토큰 임베딩의 그래디언트를 0으로 설정해야 합니다:

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

스크립트 실행

모든 변경 사항을 완료했거나 기본 구성으로 진행하는 데 문제가 없다면, 학습 스크립트를 실행할 준비가 된 것입니다! 🚀

이 가이드에서는 고양이 이미지 예제를 다운로드하여 사용합니다. 원하시는 경우, 직접 데이터셋을 생성하여 사용할 수도 있습니다(학습용 데이터셋 만들기 가이드 참고).

환경 변수 MODEL_NAME을 Hub의 모델 ID 또는 로컬 모델 경로로, INSTANCE_DIR을 방금 고양이 이미지를 다운로드한 경로로, OUTPUT_DIR을 모델을 저장할 위치로 설정합니다. 새로 학습된 임베딩을 연결할 특수 단어로 <new1>를 사용합니다. 스크립트는 모델 체크포인트와 pytorch_custom_diffusion_weights.bin 파일을 저장소에 생성하고 저장합니다.

Weights and Biases로 학습 진행 상황을 모니터링하려면 학습 명령에 --report_to=wandb 매개변수를 추가하고 --validation_prompt로 검증 프롬프트를 지정하세요. 이는 디버깅 및 중간 결과 저장에 유용합니다.

인간 얼굴로 학습하는 경우, Custom Diffusion 팀은 다음 매개변수가 잘 작동한다고 발견했습니다:

--learning_rate=5e-6

--max_train_steps은 1000에서 2000 사이의 값으로 설정할 수 있습니다

--freeze_model=crossattn

학습에 최소 15~20개의 이미지를 사용하세요

단일 개념

여러 개념

학습이 완료되면, 새로 만든 Custom Diffusion 모델을 추론에 사용할 수 있습니다.

단일 개념

여러 개념

다음 단계

Custom Diffusion으로 모델을 학습시키신 것을 축하드립니다! 🎉 더 자세히 알아보려면:

Custom Diffusion 팀의 실험 결과에 대한 자세한 내용은 텍스트-이미지 확산의 다중 개념 커스터마이징 블로그 포스트를 참고하세요.

Update on GitHub

Diffusers

Custom Diffusion

스크립트 매개변수

사전 보존 손실

정규화

학습 스크립트

스크립트 실행

다음 단계