Ddpm Fewshot Anime Face

🚀 [擴散模型課程第2單元示例微調模型]

本項目是擴散模型課程 🧨 第2單元的示例微調模型。通過使用 huggan/few-shot-anime-face 數據集對 google/ddpm-celebahq-256 模型進行微調，旨在讓用戶能夠生成卡通角色頭像。

🚀 快速開始

生成圖像

以下代碼展示瞭如何使用微調後的模型生成卡通角色頭像：

from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('your_hub_username/ddpm-fewshot_anime_face')
image = pipeline().images[0]
image.show()  # To display the generated image

微調模型

如果你想對模型進行微調，可參考以下代碼。此代碼適用於數據量較少的情況：

from torchvision import transforms
import torch
from diffusers import DDPMPipeline,DDIMScheduler
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm_notebook
import torch.nn.functional as F
from matplotlib import pyplot as plt

path = "/Volumes/mac1/pytorchlearning/diff_learning/ddpm-fewshot_anime_face" #改為model_id
image_pipe = DDPMPipeline.from_pretrained(path)
scheduler = DDIMScheduler.from_pretrained(path)
dataset = load_dataset("huggan/few-shot-anime-face",split="train")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
image_pipe.to(device)

model_name = "ddpm-fewshot_anime_face"
image_size = 256
batch_size = 8
preprocess = transforms.Compose([
    transforms.Resize(image_size),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.5],[0.5])
])

def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images":images}

dataset.set_transform(transform)
train_loader = DataLoader(dataset,batch_size,shuffle=True)

num_epochs = 20
lr = 1e-5
grad_accumulation_steps = 2
optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)
losses = []
for epoch in range(num_epochs):
    for step,batch in tqdm_notebook(enumerate(train_loader), total=len(train_loader)):
        clean_images = batch["images"].to(device)
        noise = torch.randn(clean_images.shape).to(device)
        bs = clean_images.shape[0]
        
        time_steps = torch.randint(0,image_pipe.scheduler.num_train_timesteps,(bs,),device=device).long()
        noise_image = image_pipe.scheduler.add_noise(clean_images,noise,time_steps)
        noise_pred = image_pipe.unet(noise_image,time_steps,return_dict=False)[0]
        loss = F.mse_loss(noise_pred,noise_image)
        
        losses.append(loss.item())
        loss.backward()
        if (step + 1) % grad_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

    print(f"Epoch {epoch} average loss: {sum(losses[-len(train_loader):])/len(train_loader)}")
    image_pipe.save_pretrained(model_name)

plt.plot(losses)

📄 許可證

本項目採用 MIT 許可證。