Tinh chỉnh Stable Diffusion XL với DreamBooth và LoRA

Trong hướng dẫn này, chúng ta sẽ tìm hiểu về Stable Diffusion XL và DreamBooth, cũng như cách truy cập mô hình tạo ảnh bằng thư viện diffusers. Ngoài ra, chúng ta sẽ học cách tinh chỉnh mô hình trên ảnh cá nhân và đánh giá hiệu năng của nó. Nếu bạn mới làm quen với AI, chúng tôi khuyên bạn nên học khóa AI Fundamentals để nắm kiến thức nhập môn. 

Tìm hiểu Stable Diffusion XL

Nhóm Stability AI đã phát hành Stable Diffusion XL (SDXL) 1.0, đánh dấu bước tiến tiếp theo của các mô hình AI chuyển văn bản thành hình ảnh. Mô hình nguồn mở này được xây dựng từ phiên bản SDXL 0.9 (chỉ dành cho nghiên cứu) trước đó để trở thành mô hình tạo ảnh công khai mạnh mẽ nhất hiện có.

Cả phân tích định lượng và đánh giá định tính từ con người qua nhiều tuần thử nghiệm đều cho thấy SDXL có khả năng tạo ra những hình ảnh chất lượng cao nhất và được ưa thích hơn so với các mô hình nguồn mở khác.

Hình ảnh từ arxiv.org

Chất lượng cao này đạt được nhờ sử dụng tập hợp hai mô hình — một bộ tạo cơ sở 3,5 tỷ tham số và một bộ tinh chỉnh 6,6 tỷ tham số. Quy trình kép này tối đa hóa chất lượng ảnh đồng thời vẫn đủ hiệu quả để chạy trên GPU dành cho người dùng phổ thông.

Với SDXL 1.0, người dùng không còn cần các prompt dài và phức tạp để tạo ra hình ảnh ấn tượng. “Trí thông minh” của mô hình cho phép tạo ảnh tinh xảo chỉ từ vài từ.

Việc tinh chỉnh SDXL trên các tập dữ liệu và tác vụ tùy biến đã trở nên dễ dàng hơn. Nó cho phép chúng ta kiểm soát chi tiết về bố cục, phong cách và thành phần.

DreamBooth là gì?

DreamBooth, được giới thiệu vào năm 2022 bởi nhóm nghiên cứu của Google, là một bước tiến quan trọng trong lĩnh vực AI sinh tạo, đặc biệt với các mô hình chuyển văn bản thành hình ảnh như Stable Diffusion.

Nó được gọi là DreamBooth vì, theo lời các nhà nghiên cứu Google:

"Giống như một buồng chụp ảnh nhưng ghi lại chủ thể theo cách cho phép tổng hợp họ ở bất cứ nơi nào mà giấc mơ của bạn đưa tới."

Hình ảnh từ DreamBooth

DreamBooth cho phép bạn đưa một chủ thể tùy chỉnh cụ thể vào mô hình, sau khi tinh chỉnh, mô hình sẽ chuyên biệt hóa việc thể hiện chủ thể đó theo nhiều cách khác nhau. Theo nghĩa này, nó mở ra khả năng tạo một bộ tạo ảnh của riêng bạn tập trung vào một người, nhân vật, đối tượng hoặc bối cảnh cụ thể.

DreamBooth chỉ cần một vài (thường 3–5) ảnh của chủ thể để huấn luyện mô hình hiệu quả. Sau khi huấn luyện, mô hình có thể đặt chủ thể trong vô vàn bối cảnh, cảnh và tư thế — chỉ bị giới hạn bởi trí tưởng tượng của người dùng.

Trường hợp sử dụng DreamBooth

Tinh chỉnh mô hình tạo ảnh với DreamBooth có thể hữu ích cho nhiều lĩnh vực. Nó mang lại nhiều tự do để thử nghiệm và tạo ảnh chất lượng cao mà không cần kỹ năng Photoshop. Dưới đây là một số ví dụ về cách nó có thể dễ dàng mang lại giá trị cho môi trường làm việc.

1. Các ngành sáng tạo như thiết kế đồ họa, quảng cáo và giải trí có thể hưởng lợi lớn từ DreamBooth, vì nó mang đến mức độ tùy biến và độc đáo cao trong việc tạo nội dung hình ảnh.
2. DreamBooth cho phép cá nhân hóa bằng cách tạo ra các kịch bản khó tái hiện trong đời thực hoặc hoàn toàn hư cấu.
3. Nó cũng có thể được áp dụng trong giáo dục và nghiên cứu nơi hình ảnh trực quan rất quan trọng, chẳng hạn dùng để tạo nội dung học tập cá nhân hóa hoặc phục vụ nghiên cứu.

Trong các phần tiếp theo, chúng ta sẽ đi sâu vào quy trình truy cập và tinh chỉnh mô hình SDXL trên một tập dữ liệu tùy chỉnh bằng GPU miễn phí của Kaggle.

Truy cập Stable Diffusion XL

Chúng ta có thể thử bản demo Stable Diffusion XL Spaces trên Hugging Face, nơi nhanh chóng tạo ra bốn bức ảnh dựa trên đầu vào của bạn. Hãy thử nghiệm trước khi quyết định xem nó có phù hợp với ứng dụng của bạn không.

Hình ảnh từ Stable Diffusion XL trên TPUv5e

Cách khác là dùng thư viện Python diffusers để tạo ảnh bằng prompt tùy chỉnh.

Thiết lập

Trước khi chạy mã, hãy đảm bảo bạn đang dùng máy có GPU hỗ trợ CUDA.

!nvidia-smi

Sau đó, cài đặt gói diffusers bằng PIP.

%pip install --upgrade diffusers[torch] -q

Tải mô hình cơ sở và bộ giải mã VAE

Chúng ta sẽ tải bộ giải mã VAE tùy chỉnh đã được chỉnh sửa để chạy với độ chính xác fp16 mà không sinh NaN.

Sau đó, chúng ta sẽ xây dựng pipeline bằng cách tải mô hình SDXL cơ sở với độ chính xác fp16 và tích hợp bộ giải mã VAE vào quy trình của diffusers.

Lẽ ra có thể tải trực tiếp mô hình cơ sở mà không cần “fp16”, nhưng điều đó sẽ gây vấn đề bộ nhớ GPU. Vì vậy, để chạy trên GPU của Kaggle, Colab và Laptop, bạn cần truy cập biến thể fp16 của mô hình.

from diffusers import DiffusionPipeline, AutoencoderKL
import torch

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", 
    torch_dtype=torch.float16
)
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
pipe.to("cuda");

Chạy pipeline tạo ảnh

Để tạo ảnh, bạn cần cung cấp một prompt đơn giản, số bước suy luận và số ảnh trên mỗi prompt.

Trong ví dụ này, chúng ta sẽ tạo bốn ảnh từ một prompt.

prompt = "A man in a spacesuit is running a marathon in the jungle."

image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt = 4)

Để hiển thị ảnh theo dạng lưới, ta có thể viết một hàm Python đơn giản để đổi kích thước ảnh và sắp xếp chúng theo lưới.

from PIL import Image

def image_grid(imgs, rows, cols, resize=256):
    assert len(imgs) == rows * cols

    if resize is not None:
        imgs = [img.resize((resize, resize)) for img in imgs]

    w, h = imgs[0].size
    grid_w, grid_h = cols * w, rows * h
    grid = Image.new("RGB", size=(grid_w, grid_h))

    for i, img in enumerate(imgs):
        x = i % cols * w
        y = i // cols * h
        grid.paste(img, box=(x, y))

    return grid

Chúng ta sẽ hiển thị bốn ảnh trong lưới 2x2 để dễ so sánh đầu ra của mô hình.

image_grid(image.images, 2, 2)

Kết quả thật ấn tượng. Rõ ràng, SDXL tốt hơn rất nhiều so với Stable Diffusion 1.6.

Hãy viết một prompt khác và tạo ba bức ảnh về một chú khỉ chơi pháo hoa.

prompt = "A playful monkey handling fireworks with ease."
image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt = 3)
image_grid(image.images, 1, 3)

Tìm hiểu cách truy cập mô hình Stable Diffusion trực tuyến và cục bộ bằng cách làm theo hướng dẫn Cách chạy Stable Diffusion.

Cải thiện kết quả với Refiner

Để tiếp tục cải thiện chất lượng ảnh và độ chính xác của mô hình, chúng ta sẽ dùng Refiner.

1. Tải SDXL refiner 1.0 bằng diffusion pipeline.
2. Tạo ảnh bằng mô hình SDXL cơ sở. Chúng ta sẽ đổi kiểu đầu ra thành “latent” để nhận biểu diễn tiềm ẩn của đầu vào thay vì ảnh đã tái tạo.
3. Cung cấp cho refiner prompt và biểu diễn tiềm ẩn của ảnh đã tạo.

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)
refiner.to("cuda")

n_steps = 40
high_noise_frac = 0.7

image = pipe(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
).images

image = refiner(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=image,
).images[0]

image

Ảnh sau tinh chỉnh có chất lượng cao và bám sát prompt.

Tìm hiểu cách tạo ảnh chân thực bằng Python với các mô hình khuếch tán hiện đại qua buổi code-along Generating Photorealistic Images using AI with Diffusers in Python.

Tinh chỉnh SDXL bằng AutoTrain Advanced

Theo thời gian, việc tinh chỉnh SDXL đã trở nên dễ dàng hơn nhiều. Nhờ AutoTrain Advance, giờ đây chúng ta có thể tinh chỉnh mô hình chỉ với một script Python. Thư viện Học máy Tự động này được thiết kế để huấn luyện và triển khai các mô hình máy học hiện đại với lượng mã tối thiểu.

Bạn có thể cài đặt gói Python bằng PIP:

%pip install -U autotrain-advanced

Nếu bạn quan tâm đến việc tinh chỉnh mô hình SDXL bằng accelerate và transformers, hãy xem Colab Notebook SDXL DreamBoot LoRA. Notebook này hơi lỗi thời và bạn có thể phải chỉnh sửa để chạy ổn định.

Thiết lập

Trước khi chạy script DreamBooth, chúng ta nên thiết lập một số biến sẽ dùng để chạy script.

1. Đặt tên dự án.
2. Vị trí mô hình cơ sở.
3. Thư mục chứa dữ liệu.
4. Repo ID để đẩy mô hình lên Hugging Face.

PROJECT_NAME = "Dreambooth_SDXL"
MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
DATA_DIR = "/kaggle/input/abids-photos"
REPO_ID = "kingabzpro/sdxl-lora-abid"

Tiếp theo, bạn sẽ thiết lập token Hugging Face bằng tính năng Secret của Kaggle. Tìm hiểu cách thiết lập secrets bằng cách đọc Hướng dẫn Mistral 7B: Từng bước sử dụng Mistral 7B.

from kaggle_secrets import UserSecretsClient
user_secrets = UserSecretsClient()
HF_TOKEN = user_secrets.get_secret("HUGGINGFACE_TOKEN")

Tạo tập dữ liệu riêng tư trên Kaggle

Để tạo tập dữ liệu riêng tư gồm ảnh selfie của bạn trên Kaggle, bạn cần bấm nút tải lên ở bảng điều khiển bên phải của notebook.

Thêm tiêu đề cho tập dữ liệu và tải lên tệp zip chứa năm bức ảnh selfie chụp ở các địa điểm khác nhau. Ở đây là ảnh của tác giả.

Để xem cả năm ảnh trong Jupyter Notebook, chúng ta sẽ tải ảnh bằng PIL và đưa vào hàm image_grid.

from PIL import Image

def image_grid(imgs, rows, cols, resize=256):
    assert len(imgs) == rows * cols

    if resize is not None:
        imgs = [img.resize((resize, resize)) for img in imgs]

    w, h = imgs[0].size
    grid_w, grid_h = cols * w, rows * h
    grid = Image.new("RGB", size=(grid_w, grid_h))

    for i, img in enumerate(imgs):
        x = i % cols * w
        y = i // cols * h
        grid.paste(img, box=(x, y))

    return grid

import glob

Chúng ta đã hiển thị năm ảnh chất lượng cao của tác giả tại các địa điểm khác nhau. Hãy đảm bảo tất cả ảnh có chất lượng cao và kích thước tương đồng.

imgs = [Image.open(path) for path in glob.glob("/kaggle/input/abids-photos/*.jpg")]
image_grid(imgs, 1, 5)

Script huấn luyện AutoTrain DreamBooth

Đã đến lúc tinh chỉnh mô hình của chúng ta trên năm bức ảnh này bằng script AutoTrain DreamBooth.

1. Cung cấp mô hình cơ sở, tên dự án và thư mục dữ liệu.
2. Đảm bảo instance prompt là duy nhất, mô tả người hoặc đối tượng. Dùng tên đầy đủ của người, "Abid Ali Awan", và mô tả chi tiết bức ảnh.
3. Giữ các tham số còn lại như mặc định để đạt hiệu năng tốt nhất.
4. Cung cấp access token của Hugging Face và repository ID để tải adapter LoRA của mô hình lên Hugging Face Hub.

!autotrain dreambooth \
--model $MODEL_NAME \
--project-name $PROJECT_NAME \
--image-path $DATA_DIR \
--prompt "A photo of Abid Ali Awan wearing casual clothes, taking a selfie, and smiling." \
--resolution 1024 \
--batch-size 1 \
--num-steps 500 \
--gradient-accumulation 4 \
--lr 1e-4 \
--fp16 \
--gradient-checkpointing \
--push-to-hub \
--token $HF_TOKEN \
--repo-id $REPO_ID

Mất 2 giờ 15 phút để huấn luyện mô hình và đẩy trọng số LoRA lên Hugging Face.

LoRA, viết tắt của Low-Rank Adaptation, là một kỹ thuật bổ sung các lớp có thể huấn luyện nhỏ vào mô hình hiện có mà không sửa đổi trọng số gốc.

LoRA hữu ích vì cho phép đưa các khái niệm mới, như phong cách nghệ thuật, nhân vật hoặc chủ đề, vào mô hình mà không cần nhiều tính toán hay bộ nhớ.

Steps: 100%|█████████| 500/500 [2:15:11<00:00, 14.49s/it, loss=0.119, lr=0.0001]Model weights saved in Dreambooth_SDXL/pytorch_lora_weights.safetensors
Steps: 100%|█████████| 500/500 [2:15:11<00:00, 16.22s/it, loss=0.119, lr=0.0001]
pytorch_lora_weights.safetensors: 100%|████| 23.4M/23.4M [00:01<00:00, 11.9MB/s]

Thay vì lưu một tệp mô hình 7GB, chúng ta sẽ lưu một LoRA adapter dung lượng 23,4MB để gắn vào mô hình cơ sở trong quá trình suy luận.

Hình ảnh từ kingabzpro/sdxl-lora-abid

Suy luận với mô hình SDXL đã tinh chỉnh

Chúng ta sẽ tạo một diffusion pipeline để sinh ảnh, sử dụng bộ giải mã VAE tùy chỉnh và mô hình SDXL cơ sở với độ chính xác FP16.

Tiếp theo, chúng ta sẽ tải trọng số LoRA và thêm vào mô hình cơ sở bằng repo ID trên Hugging Face.

Cuối cùng, chúng ta sẽ chạy pipeline với prompt và tạo ba bức ảnh.

from diffusers import DiffusionPipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline

import torch

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", 
    torch_dtype=torch.float16
)
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
pipe.to("cuda");
pipe.load_lora_weights(REPO_ID, weight_name="pytorch_lora_weights.safetensors")


prompt = "A photo of Abid Ali Awan participating in a marathon."

image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt = 3)
image_grid(image.images, 1, 3)

Kết quả rất ấn tượng. Chúng ta đã tinh chỉnh thành công mô hình để tạo ảnh của tác giả “Abid Ali Awan”.

Hãy thử lại với một prompt khác.

prompt = "A photo of Abid Ali Awan giving a press conference."

image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt = 3)
image_grid(image.images, 1, 3)

Kết quả một lần nữa rất xuất sắc.

Sử dụng Refiner

Trong phần này, chúng ta sẽ cải thiện ảnh đã tạo bằng SDXL Refiner 1.0.

Trước tiên, chúng ta sẽ tạo ảnh bằng pipeline Stable Diffusion và đặt seed thủ công để đảm bảo khả năng tái lập.

prompt = "A photo of Abid Ali Awan wearing a business suit at an office meeting."

seed = 65
generator = torch.Generator("cuda").manual_seed(seed)
image = pipe(prompt=prompt, num_inference_steps=25, generator=generator).images[0]
image.resize((300, 300))

Tiếp theo, chúng ta sẽ tải SDXL refiner và chạy pipeline refiner với cùng prompt, generator và ảnh đã tạo.

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)
refiner.to("cuda");

image = refiner(prompt=prompt, num_inference_steps=25, generator=generator, image=image)
image.images[0].resize((300, 300))

Refiner đã tăng cường chất lượng ảnh, nhưng kết quả không còn giống ảnh gốc của tác giả. Thay vào đó, nó tạo ra hình ảnh một doanh nhân phổ thông ở Nam Á. Có lẽ, với các mô hình đã tinh chỉnh, tốt hơn là nên tránh dùng refiner.

Sau khi tinh chỉnh mô hình SDXL, bước tiếp theo là xây dựng ứng dụng AI sinh tạo. Để lấy cảm hứng cho dự án, hãy xem blog 5 dự án xây dựng với mô hình sinh tạo và công cụ nguồn mở.

Kết luận

Trong hướng dẫn này, chúng ta đã tìm hiểu về mô hình Stable Diffusion XL và kỹ thuật DreamBooth. Ngoài ra, chúng ta đã học cách tinh chỉnh Stable Diffusion XL bằng script DreamBooth của AutoTrain Advance để tạo ảnh cá nhân hóa.

Sau khi tải mô hình SDXL cơ sở, chúng ta đã huấn luyện một LoRA adapter gọn nhẹ chỉ với năm bức ảnh của tác giả. Sau đó, trọng số LoRA được gắn vào SDXL cơ sở, cho phép mô hình tạo ảnh chất lượng cao của tác giả trong nhiều kịch bản tưởng tượng khác nhau, đồng thời vẫn giữ được năng lực của mô hình nhiều tỷ tham số đầy đủ.

Mặc dù SDXL Refiner cải thiện chất lượng ảnh, nó hiệu quả nhất với các prompt chung chung và dường như ghi đè một phần tính cá nhân hóa đạt được qua tinh chỉnh. Dù vậy, hướng dẫn đã cho thấy việc tinh chỉnh mô hình SDXL bằng AutoTrain Advance với chỉ vài bức ảnh là rất đơn giản.

Nếu bạn mới bước vào thế giới AI và muốn hiểu vì sao nó được quan tâm, chúng tôi khuyên bạn học khóa AI Fundamentals. Khóa học sẽ giúp bạn tìm hiểu về AI sinh tạo và các mô hình ngôn ngữ phức tạp. Bên cạnh đó, bạn có thể khởi đầu sự nghiệp kỹ sư AI bằng cách tham gia lộ trình nghề nghiệp Machine Learning Scientist with Python. Chương trình trọn vẹn này sẽ giúp bạn học xử lý ngôn ngữ tự nhiên và xử lý ảnh, cũng như các gói máy học Python phổ biến như scikit-learn, PySpark và Keras.