Diffusers documentation
Latent Consistency Distillation
Latent Consistency Distillation
Latent Consistency Models (LCMs)는 단 몇 단계 만에 고품질 이미지를 생성할 수 있어, 많은 파이프라인이 최소 25단계 이상을 요구하는 점을 고려할 때 큰 진전을 의미합니다. LCMs는 어떤 Stable Diffusion 모델에도 latent consistency distillation 방법을 적용하여 생성됩니다. 이 방법은 latent space에 one-stage guided distillation을 적용하고, skipping-step 방법을 결합하여 timesteps를 일관되게 건너뛰어 증류 과정을 가속화하는 방식으로 작동합니다(자세한 내용은 논문의 4.1, 4.2, 4.3절을 참조하세요).
vRAM이 제한된 GPU에서 학습하는 경우, gradient_checkpointing, gradient_accumulation_steps, mixed_precision를 활성화하여 메모리 사용량을 줄이고 학습 속도를 높일 수 있습니다. xFormers와 bitsandbytes’ 8-bit 옵티마이저를 사용하여 메모리 효율적인 어텐션을 활성화하면 메모리 사용량을 더욱 줄일 수 있습니다.
이 가이드에서는 train_lcm_distill_sd_wds.py 스크립트를 살펴보고, 이를 더 잘 이해하며 사용자의 용도에 맞게 수정하는 방법을 안내합니다.
…
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .이후 학습 스크립트가 포함된 예제 폴더로 이동하여 사용하는 스크립트에 필요한 의존성을 설치합니다:
cd examples/consistency_distillation
pip install -r requirements.txt🤗 Accelerate는 여러 GPU/TPU 또는 혼합 정밀도로 학습하는 것을 돕는 라이브러리입니다. 하드웨어와 환경에 따라 학습 설정을 자동으로 구성합니다. 자세한 내용은 🤗 Accelerate Quick tour를 참고하세요.
🤗 Accelerate 환경을 초기화합니다(torch.compile을 활성화하면 학습 속도가 크게 빨라집니다):
accelerate config
어떤 구성도 선택하지 않고 기본 🤗 Accelerate 환경을 설정하려면:
accelerate config default
또는 노트북과 같이 인터랙티브 셸을 지원하지 않는 환경에서는 다음을 사용할 수 있습니다:
from accelerate.utils import write_basic_config
write_basic_config()마지막으로, 자체 데이터셋으로 모델을 학습하려면 학습용 데이터셋 만들기 가이드를 참고하여 학습 스크립트와 호환되는 데이터셋을 만드는 방법을 배워보세요.
스크립트 매개변수
다음 섹션에서는 스크립트를 수정하는 방법을 이해하는 데 중요한 학습 스크립트의 일부 부분을 강조합니다. 그러나 스크립트의 모든 측면을 자세히 다루지는 않습니다. 더 자세히 알고 싶으시다면 스크립트를 읽어보시고, 질문이나 우려 사항이 있으시면 알려주시기 바랍니다.
학습 스크립트는 학습 실행을 사용자 정의하는 데 도움이 되는 많은 매개변수를 제공합니다. 모든 매개변수와 그 설명은 parse_args() 함수에서 찾을 수 있습니다. 이 함수는 학습 배치 크기나 학습률과 같이 각 매개변수에 대한 기본 값을 제공하지만, 원하시는 경우 학습 명령에서 직접 값을 설정할 수도 있습니다.
예를 들어, fp16 형식을 사용하여 혼합 정밀도로 학습 속도를 높이려면 학습 명령에 --mixed_precision 매개변수를 추가합니다:
accelerate launch train_lcm_distill_sd_wds.py \
--mixed_precision="fp16"대부분의 파라미터는 텍스트-투-이미지 학습 가이드의 파라미터와 동일하므로, 이 가이드에서는 잠재적 일관성 증류에 관련되는 파라미터에 초점을 맞춥니다.
--pretrained_teacher_model: teacher model로 사용할 사전 학습된 latent diffusion model의 경로--pretrained_vae_model_name_or_path: 사전 학습된 VAE의 경로; SDXL VAE는 수치적 불안정성이 있는 것으로 알려져 있으므로, 이 매개변수를 통해 fp16에서 작동하는 madebyollin의 VAE와 같은 대안 VAE를 지정할 수 있습니다--w_min및--w_max: guidance scale 샘플링을 위한 guidance scale의 최소값과 최대값--num_ddim_timesteps: DDIM 샘플링에 사용할 타임스텝 수--loss_type: 잠재적 일관성 증류(latent consistency distillation)에 대해 계산할 손실(L2 또는 Huber)의 유형; Huber 손실은 이상치에 대해 더 강건하여 일반적으로 선호됩니다--huber_c: Huber 손실 함수의 파라미터
학습 스크립트
학습 스크립트는 이미지 전처리 및 학습 데이터셋 생성을 위한 데이터셋 클래스 Text2ImageDataset를 생성하는 것으로 시작합니다.
def transform(example):
image = example["image"]
image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)
c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
image = TF.crop(image, c_top, c_left, resolution, resolution)
image = TF.to_tensor(image)
image = TF.normalize(image, [0.5], [0.5])
example["image"] = image
return example클라우드에 저장된 대규모 데이터셋의 읽기 및 쓰기 성능을 향상시키기 위해, 이 스크립트는 WebDataset 형식을 사용하여 변환을 적용하고 학습용 데이터셋 및 dataloader를 생성하는 전처리 파이프라인을 만듭니다. 전체 데이터셋을 먼저 다운로드할 필요 없이 이미지가 처리되어 학습 루프에 공급됩니다.
processing_pipeline = [
wds.decode("pil", handler=wds.ignore_and_continue),
wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
wds.map(filter_keys({"image", "text"})),
wds.map(transform),
wds.to_tuple("image", "text"),
]main() 함수에서는 노이즈 스케줄러, 토크나이저, 텍스트 인코더, VAE와 같은 필요한 모든 구성 요소를 로드합니다. 또한 여기서 teacher UNet을 로드한 후, teacher UNet을 기반으로 student UNet을 생성할 수 있습니다. student UNet은 학습 중에 옵티마이저에 의해 업데이트됩니다.
teacher_unet = UNet2DConditionModel.from_pretrained(
args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)
unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()이제 UNet의 매개변수를 업데이트하기 위한 optimizer를 생성할 수 있습니다:
optimizer = optimizer_class(
unet.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)dataset를 생성합니다:
dataset = Text2ImageDataset(
train_shards_path_or_url=args.train_shards_path_or_url,
num_train_examples=args.max_train_samples,
per_gpu_batch_size=args.train_batch_size,
global_batch_size=args.train_batch_size * accelerator.num_processes,
num_workers=args.dataloader_num_workers,
resolution=args.resolution,
shuffle_buffer_size=1000,
pin_memory=True,
persistent_workers=True,
)
train_dataloader = dataset.train_dataloader다음으로 training loop를 설정하고 latent consistency distillation 방법을 구현할 준비가 되었습니다(자세한 내용은 논문 Algorithm 1을 참고하세요). 이 스크립트 섹션에서는 latents에 노이즈를 추가하고, 샘플링을 수행하여 guidance scale 임베딩을 생성하며, 노이즈로부터 원본 이미지를 예측하는 작업을 처리합니다.
pred_x_0 = predicted_origin(
noise_pred,
start_timesteps,
noisy_model_input,
noise_scheduler.config.prediction_type,
alpha_schedule,
sigma_schedule,
)
model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0다음으로 teacher model predictions와 LCM predictions를 가져와 손실을 계산한 뒤 LCM으로 역전파합니다.
if args.loss_type == "l2":
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
loss = torch.mean(
torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
)If you want to learn more about how the training loop works, check out the Understanding pipelines, models and schedulers tutorial which breaks down the basic pattern of the denoising process.
스크립트 실행
이제 학습 스크립트를 실행하여 증류를 시작할 준비가 되었습니다!
이 가이드에서는 여기에 저장된 Hub의 Conceptual Captions 12M 데이터셋 경로를 지정하기 위해 --train_shards_path_or_url을 사용합니다. MODEL_DIR 환경 변수를 teacher 모델 이름으로, OUTPUT_DIR을 모델을 저장할 위치로 설정하세요.
export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"
accelerate launch train_lcm_distill_sd_wds.py \
--pretrained_teacher_model=$MODEL_DIR \
--output_dir=$OUTPUT_DIR \
--mixed_precision=fp16 \
--resolution=512 \
--learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
--max_train_steps=1000 \
--max_train_samples=4000000 \
--dataloader_num_workers=8 \
--train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
--validation_steps=200 \
--checkpointing_steps=200 --checkpoints_total_limit=10 \
--train_batch_size=12 \
--gradient_checkpointing --enable_xformers_memory_efficient_attention \
--gradient_accumulation_steps=1 \
--use_8bit_adam \
--resume_from_checkpoint=latest \
--report_to=wandb \
--seed=453645634 \
--push_to_hub학습이 완료되면 새로운 LCM을 추론에 사용할 수 있습니다.
from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch
unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")
pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda") # or "mps", "xpu", "cpu"
prompt = "sushi rolls in the form of panda heads, sushi platter"
image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]LoRA
LoRA는 학습 가능한 매개변수의 수를 크게 줄이는 학습 기법입니다. 그 결과, 학습 속도가 빨라지고, 생성된 가중치가 훨씬 작아(~100MB) 저장하기가 쉬워집니다. LoRA로 학습하려면 train_lcm_distill_lora_sd_wds.py 또는 train_lcm_distill_lora_sdxl.wds.py 스크립트를 사용하세요.
LoRA 학습 스크립트는 LoRA 학습 가이드에서 더 자세히 설명되어 있습니다.
Stable Diffusion XL
Stable Diffusion XL (SDXL)은 고해상도 이미지를 생성하는 강력한 텍스트-투-이미지 모델이며, 아키텍처에 두 번째 텍스트 인코더를 추가했습니다. train_lcm_distill_sdxl_wds.py 스크립트를 사용하여 LoRA로 SDXL 모델을 학습하세요.
SDXL 학습 스크립트는 SDXL 학습 가이드에서 더 자세히 설명되어 있습니다.
다음 단계
LCM 모델을 증류하신 것을 축하드립니다! LCM에 대해 더 자세히 알고 싶으시다면 다음 내용이 도움이 될 수 있습니다:
- 텍스트-이미지, 이미지-이미지, LoRA 체크포인트와 함께 추론을 위한 LCM 사용법을 학습해 보세요.
- SDXL LCM-LoRA의 초고속 추론, 품질 비교, 벤치마크 등에 대해 자세히 알아보려면 Latent Consistency LoRAs를 활용한 4단계 SDXL 블로그 포스트를 읽어보세요.