stable diffusion 训练的伪代码

import torch
from diffusers import StableDiffusionPipeline
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os

# 1. 初始化模型和数据
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base").to("cuda")

class CustomDataset(Dataset):
    def __init__(self, data_dir, tokenizer):
        self.data_dir = data_dir
        self.tokenizer = tokenizer
        self.data = self.load_data()

    def load_data(self):
        data = []
        for file in os.listdir(self.data_dir):
           image_path = os.path.join(self.data_dir, file)
           text_path = os.path.join(self.data_dir, file.replace(".jpg", ".txt").replace(".png", ".txt"))
           if os.path.exists(text_path):
               with open(text_path, "r") as f:
                    text = f.read()
            data.append((image_path, text))
        return data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image_path, text = self.data[idx]
        image = Image.open(image_path).convert("RGB")
        inputs = self.tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=77)
        return {"image": image, "input_ids": inputs["input_ids"].flatten()}

dataset = CustomDataset(data_dir="path_to_your_data", tokenizer=tokenizer)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

# 2. 数据预处理
def preprocess_image(image):
    image = image.resize((512, 512))
    image = torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).to("cuda", dtype=torch.float16) / 127.5 - 1.0
    return image

# 3. 训练循环
optimizer = torch.optim.AdamW(pipe.unet.parameters(), lr=1e-5)
for epoch in range(10):
    for batch in dataloader:
        image = preprocess_image(batch["image"])
        input_ids = batch["input_ids"].to("cuda")
  

        # 计算损失
        noise = torch.randn_like(image)
        timesteps = torch.randint(0, pipe.scheduler.num_train_timesteps, (image.shape[0],)).to("cuda")
        noisy_image = pipe.scheduler.add_noise(image, noise, timesteps)
        noise_pred = pipe.unet(noisy_image, timesteps, context=model(input_ids, attention_mask=attention_mask)[0]).sample
        loss = torch.nn.functional.mse_loss(noise_pred, noise)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 4. 保存模型
    pipe.save_pretrained(f"custom_model_epoch_{epoch}")