Memahami pipeline, model, dan scheduler

🧨 Diffusers dirancang sebagai kotak alat yang ramah pengguna dan fleksibel untuk membangun sistem difusi yang disesuaikan dengan kebutuhan Anda. Inti dari kotak alat ini adalah model dan scheduler. Meskipun DiffusionPipeline mengemas komponen-komponen ini bersama untuk kenyamanan, Anda juga dapat memisahkan pipeline tersebut dan menggunakan model serta scheduler secara terpisah untuk membuat sistem difusi baru.

Dalam tutorial ini, Anda akan belajar cara menggunakan model dan scheduler untuk merakit sistem difusi untuk inferensi, dimulai dari pipeline dasar dan kemudian beralih ke pipeline Stable Diffusion.

Membongkar pipeline dasar

Pipeline adalah cara cepat dan mudah untuk menjalankan model untuk inferensi, yang membutuhkan tidak lebih dari empat baris kode untuk menghasilkan gambar:

>>> from diffusers import DDPMPipeline

>>> ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
>>> image = ddpm(num_inference_steps=25).images[0]
>>> image

Itu sangat mudah, tetapi bagaimana pipeline melakukannya? Mari kita bedah pipeline dan lihat apa yang terjadi di balik layar.

Pada contoh di atas, pipeline berisi model UNet2DModel dan DDPMScheduler. Pipeline menghilangkan noise pada gambar dengan mengambil noise acak sebesar output yang diinginkan dan melewatkan model beberapa kali. Pada setiap timestep, model memprediksi noise residual dan scheduler menggunakannya untuk memprediksi gambar yang lebih sedikit noise. Pipeline mengulangi proses ini hingga mencapai akhir dari jumlah langkah inferensi yang ditentukan.

Untuk membuat ulang pipeline dengan model dan scheduler secara terpisah, mari kita tulis proses denoising kita sendiri.

Muat model dan scheduler:

>>> from diffusers import DDPMScheduler, UNet2DModel

>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
>>> model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

Atur jumlah timestep untuk menjalankan proses denoising:

>>> scheduler.set_timesteps(50)

Menetapkan timestep scheduler membuat tensor dengan elemen yang berjarak sama, 50 dalam contoh ini. Setiap elemen sesuai dengan timestep di mana model melakukan denoising pada gambar. Saat Anda membuat loop denoising nanti, Anda akan mengiterasi tensor ini untuk melakukan denoising pada gambar:

>>> scheduler.timesteps
tensor([980, 960, 940, 920, 900, 880, 860, 840, 820, 800, 780, 760, 740, 720,
    700, 680, 660, 640, 620, 600, 580, 560, 540, 520, 500, 480, 460, 440,
    420, 400, 380, 360, 340, 320, 300, 280, 260, 240, 220, 200, 180, 160,
    140, 120, 100,  80,  60,  40,  20,   0])

Buat beberapa noise acak dengan bentuk yang sama seperti output yang diinginkan:

>>> import torch

>>> sample_size = model.config.sample_size
>>> noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")  # or "mps", "xpu", "cpu"

Sekarang tulis sebuah loop untuk mengiterasi timesteps. Pada setiap timestep, model melakukan satu pass UNet2DModel.forward() dan mengembalikan residual berisik. Metode step() dari scheduler menerima residual berisik, timestep, dan input, lalu memprediksi gambar pada timestep sebelumnya. Output ini menjadi input berikutnya untuk model dalam loop denoising, dan proses ini akan diulang hingga mencapai akhir array timesteps.

>>> input = noise

>>> for t in scheduler.timesteps:
...     with torch.no_grad():
...         noisy_residual = model(input, t).sample
...     previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
...     input = previous_noisy_sample

Ini adalah seluruh proses denoising, dan Anda dapat menggunakan pola yang sama ini untuk menulis sistem difusi apa pun.

Langkah terakhir adalah mengonversi output yang telah dinoise menjadi gambar:

>>> from PIL import Image
>>> import numpy as np

>>> image = (input / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

Pada bagian berikutnya, Anda akan menguji kemampuan Anda dan menguraikan pipeline Stable Diffusion yang lebih kompleks. Langkah-langkahnya lebih kurang sama. Anda akan menginisialisasi komponen yang diperlukan, dan mengatur jumlah timestep untuk membuat array timestep. Array timestep digunakan dalam loop denoising, dan untuk setiap elemen dalam array ini, model memprediksi gambar yang lebih sedikit berisik. Loop denoising mengiterasi timestep, dan pada setiap timestep, loop tersebut menghasilkan residual berisik dan scheduler menggunakannya untuk memprediksi gambar yang lebih sedikit berisik pada timestep sebelumnya. Proses ini diulang hingga Anda mencapai akhir array timestep.

Mari kita coba!

Membongkar pipeline Stable Diffusion

Stable Diffusion adalah model latent diffusion dari teks ke gambar. Model ini disebut model latent diffusion karena bekerja dengan representasi berdimensi lebih rendah dari gambar, bukan ruang piksel aktual, yang membuatnya lebih efisien dalam penggunaan memori. Encoder memampatkan gambar menjadi representasi yang lebih kecil, dan decoder mengubah representasi yang dipampatkan tersebut kembali menjadi gambar. Untuk model teks ke gambar, Anda memerlukan tokenizer dan encoder untuk menghasilkan embedding teks. Dari contoh sebelumnya, Anda sudah tahu bahwa Anda memerlukan model UNet dan scheduler.

Seperti yang dapat Anda lihat, ini sudah lebih kompleks daripada pipeline DDPM yang hanya berisi model UNet. Model Stable Diffusion memiliki tiga model pretrained yang terpisah.

💡 Baca blog Bagaimana cara kerja Stable Diffusion? untuk detail lebih lanjut tentang cara kerja model VAE, UNet, dan text encoder.

Sekarang setelah Anda mengetahui apa yang Anda butuhkan untuk pipeline Stable Diffusion, muat semua komponen ini dengan metode from_pretrained(). Anda dapat menemukannya di checkpoint pretrained stable-diffusion-v1-5/stable-diffusion-v1-5, dan setiap komponen disimpan dalam subfolder terpisah:

>>> from PIL import Image
>>> import torch
>>> from transformers import CLIPTextModel, CLIPTokenizer
>>> from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

>>> vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
>>> tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
>>> text_encoder = CLIPTextModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
... )
>>> unet = UNet2DConditionModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
... )

Alih-alih menggunakan PNDMScheduler bawaan, ganti dengan UniPCMultistepScheduler untuk melihat betapa mudahnya memasang scheduler yang berbeda:

>>> from diffusers import UniPCMultistepScheduler

>>> scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

Untuk mempercepat inferensi, pindahkan model ke GPU karena, berbeda dengan scheduler, model memiliki bobot yang dapat dilatih:

>>> torch_device = "cuda"  # or "mps", "xpu", "cpu"
>>> vae.to(torch_device)
>>> text_encoder.to(torch_device)
>>> unet.to(torch_device)

Buat embedding teks

Langkah selanjutnya adalah melakukan tokenisasi teks untuk menghasilkan embedding. Teks digunakan untuk mengondisikan model UNet dan mengarahkan proses difusi menuju sesuatu yang menyerupai prompt input.

💡 Parameter guidance_scale menentukan seberapa besar bobot yang diberikan pada prompt saat menghasilkan gambar.

Silakan pilih prompt apa pun yang Anda inginkan jika Anda ingin menghasilkan sesuatu yang lain!

>>> prompt = ["a photograph of an astronaut riding a horse"]
>>> height = 512  # default height of Stable Diffusion
>>> width = 512  # default width of Stable Diffusion
>>> num_inference_steps = 25  # Number of denoising steps
>>> guidance_scale = 7.5  # Scale for classifier-free guidance
>>> generator = torch.manual_seed(0)  # Seed generator to create the initial latent noise
>>> batch_size = len(prompt)

Tokenisasi teks dan generate embedding dari prompt:

>>> text_input = tokenizer(
...     prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt"
... )

>>> with torch.no_grad():
...     text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

Anda juga perlu menghasilkan unconditional text embeddings, yaitu embeddings untuk token padding. Embeddings ini harus memiliki bentuk yang sama (batch_size dan seq_length) dengan text_embeddings bersyarat:

>>> max_length = text_input.input_ids.shape[-1]
>>> uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
>>> uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

Mari kita gabungkan embedding kondisional dan non-kondisional menjadi satu batch untuk menghindari dua kali forward pass:

>>> text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

Buat noise acak

Selanjutnya, buat beberapa noise acak awal sebagai titik awal untuk proses difusi. Ini adalah representasi laten dari gambar, dan akan secara bertahap dinoised. Pada titik ini, gambar latent lebih kecil daripada ukuran gambar akhir, tetapi itu tidak masalah karena model akan mengubahnya menjadi dimensi gambar akhir 512x512 nanti.

💡 Tinggi dan lebar dibagi dengan 8 karena model vae memiliki 3 lapisan down-sampling. Anda dapat memeriksanya dengan menjalankan berikut:
2 ** (len(vae.config.block_out_channels) - 1) == 8

>>> latents = torch.randn(
...     (batch_size, unet.config.in_channels, height // 8, width // 8),
...     generator=generator,
...     device=torch_device,
... )

Hilangkan derau pada gambar

Mulai dengan menskalakan input menggunakan distribusi noise awal, sigma, yaitu nilai skala noise, yang diperlukan untuk scheduler yang ditingkatkan seperti UniPCMultistepScheduler:

>>> latents = latents * scheduler.init_noise_sigma

Langkah terakhir adalah membuat loop denoising yang akan secara progresif mengubah noise murni di latents menjadi gambar yang digambarkan oleh prompt Anda. Ingat, loop denoising perlu melakukan tiga hal:

Atur timesteps scheduler yang akan digunakan selama denoising.
Iterasi melalui langkah-langkah waktu.
Pada setiap timestep, panggil model UNet untuk memprediksi residual noise dan teruskan hasilnya ke scheduler untuk menghitung sampel berisik sebelumnya.

>>> from tqdm.auto import tqdm

>>> scheduler.set_timesteps(num_inference_steps)

>>> for t in tqdm(scheduler.timesteps):
...     # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
...     latent_model_input = torch.cat([latents] * 2)

...     latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

...     # perform guidance
...     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
...     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

...     # compute the previous noisy sample x_t -> x_t-1
...     latents = scheduler.step(noise_pred, t, latents).prev_sample

Decode gambar

Langkah terakhir adalah menggunakan vae untuk mendekode representasi laten menjadi gambar dan mendapatkan output yang didekode dengan sample:

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

Terakhir, konversi gambar menjadi PIL.Image untuk melihat gambar yang Anda hasilkan!

>>> image = (image / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

Langkah selanjutnya

Dari pipeline yang sederhana hingga yang kompleks, Anda telah melihat bahwa yang benar-benar Anda butuhkan untuk menulis sistem difusi Anda sendiri adalah loop denoising. Loop tersebut harus mengatur timesteps dari scheduler, mengiterasikannya, dan bergantian antara memanggil model UNet untuk memprediksi residual noise dan melewatkan hasilnya ke scheduler untuk menghitung sampel berisik sebelumnya.

Inilah yang sebenarnya dirancang oleh 🧨 Diffusers: membuat penulisan sistem difusi Anda sendiri menjadi intuitif dan mudah menggunakan model dan scheduler.

Untuk langkah selanjutnya, silakan:

Pelajari cara membangun dan berkontribusi pipeline ke 🧨 Diffusers. Kami tidak sabar melihat apa yang akan Anda ciptakan!
Jelajahi pipeline yang sudah ada di pustaka ini, dan lihat apakah Anda dapat membongkar dan membangun pipeline dari awal menggunakan model dan scheduler secara terpisah.

Update on GitHub