【AIGC系列】4：Stable Diffusion应用实践和代码分析 - ToB企服应用市场:ToB评测及商务社交产业平台

import torch
from diffusers import AutoencoderKL
import numpy as np
from PIL import Image
print(torch.cuda.is_available())
#加载模型: autoencoder可以通过SD权重指定subfolder来单独加载
print("Start...")
autoencoder = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
autoencoder.to("cuda", dtype=torch.float16)
print("Get weight successfully")
# 读取图像并预处理
# raw_image = Image.open("liuyifei.jpg").convert("RGB").resize((256, 256))
raw_image = Image.open("liuyifei.jpg").convert("RGB")
image = np.array(raw_image).astype(np.float32) / 127.5 - 1.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
# 压缩图像为latent并重建
with torch.inference_mode():
latent = autoencoder.encode(image.to("cuda", dtype=torch.float16)).latent_dist.sample()
rec_image = autoencoder.decode(latent).sample
rec_image = (rec_image / 2 + 0.5).clamp(0, 1)
rec_image = rec_image.cpu().permute(0, 2, 3, 1).numpy()
rec_image = (rec_image * 255).round().astype("uint8")
rec_image = Image.fromarray(rec_image[0])
rec_image.save("liuyifei_re.jpg")

复制代码

autoencoder = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse/")

复制代码

from transformers import CLIPTextModel, CLIPTokenizer
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder").to("cuda")
# text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")
tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
# tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
# 对输入的text进行tokenize，得到对应的token ids
prompt = "a photograph of an astronaut riding a horse"
text_input_ids = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).input_ids
print(f" \n\n text_input_ids: {text_input_ids} \n\n")
# 将token ids送入text model得到77x768的特征
text_embeddings = text_encoder(text_input_ids.to("cuda"))[0]
print(f" \n\n text_embeddings: {text_embeddings} \n\n")

复制代码

text_input_ids: tensor([[49406, 320, 8853, 539, 550, 18376, 6765, 320, 4558, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
49407, 49407, 49407, 49407, 49407, 49407, 49407]])
text_embeddings: tensor([[[-0.3884, 0.0229, -0.0522, ..., -0.4899, -0.3066, 0.0675],
[ 0.0290, -1.3258, 0.3085, ..., -0.5257, 0.9768, 0.6652],
[ 0.4595, 0.5617, 1.6663, ..., -1.9515, -1.2307, 0.0104],
...,
[-3.0421, -0.0656, -0.1793, ..., 0.3943, -0.0190, 0.7664],
[-3.0551, -0.1036, -0.1936, ..., 0.4236, -0.0189, 0.7575],
[-2.9854, -0.0832, -0.1715, ..., 0.4355, 0.0095, 0.7485]]],
device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

复制代码

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
import torch.nn.functional as F
# 自定义Dataset类
class CustomImageTextDataset(Dataset):
def __init__(self, image_paths, text_descriptions, transform=None):
self.image_paths = image_paths
self.text_descriptions = text_descriptions
self.transform = transform
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image_path = self.image_paths[idx]
text_description = self.text_descriptions[idx]
# 加载图像
image = Image.open(image_path).convert("RGB")
if self.transform:
image = self.transform(image)
return {
'image': image,
'text': text_description
}
# 数据准备
image_paths = ["path/to/image1.jpg", "path/to/image2.jpg"] # 替换为实际的图像路径
text_descriptions = ["description for image1", "description for image2"] # 替换为实际的文本描述
# 图像转换（预处理）
transform = transforms.Compose([
transforms.Resize((256, 256)), # 调整大小
transforms.ToTensor(), # 转换为张量
])
# 创建数据集实例
dataset = CustomImageTextDataset(image_paths=image_paths, text_descriptions=text_descriptions, transform=transform)
# 创建DataLoader
train_dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)
# 加载autoencoder
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
# 加载text encoder
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
model_config = {
"sample_size": 32,
"in_channels": 4,
"out_channels": 4,
"down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"),
"up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
"block_out_channels": (320, 640, 1280, 1280),
"layers_per_block": 2,
"cross_attention_dim": 768,
"attention_head_dim": 8,
}
# 初始化UNet
unet = UNet2DConditionModel(**model_config)
# 定义scheduler
noise_scheduler = DDPMScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
num_train_timesteps=1000
)
# 冻结vae和text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
opt = torch.optim.AdamW(unet.parameters(), lr=1e-4)
# 训练循环
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
unet.to(device)
vae.to(device)
text_encoder.to(device)
for epoch in range(10): # 假设训练10个epoch
unet.train()
for step, batch in enumerate(train_dataloader):
with torch.no_grad():
# 将image转到latent空间
latents = vae.encode(batch["image"].to(device)).latent_dist.sample()
# rescaling latents
latents = latents * vae.config.scaling_factor
# 提取text embeddings
text_input_ids = tokenizer(
batch["text"],
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).input_ids.to(device)
text_embeddings = text_encoder(text_input_ids)[0]
# 随机采样噪音
noise = torch.randn_like(latents)
bsz = latents.shape[0]
# 随机采样timestep
timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=device).long()
# 将noise添加到latent上，即扩散过程
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
# 预测noise并计算loss
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
opt.zero_grad()
loss.backward()
opt.step()
if step % 10 == 0:
print(f"Epoch {epoch}, Step {step}, Loss: {loss.item()}")
# 在训练完成后保存模型
model_save_path = 'path/to/your/unet_model.pth'
torch.save(unet.state_dict(), model_save_path)
print(f"Model has been saved to {model_save_path}")
optimizer_save_path = 'path/to/your/optimizer.pth'
torch.save(opt.state_dict(), optimizer_save_path)
print(f"Optimizer state has been saved to {optimizer_save_path}")
# 加载模型进行推理或继续训练
unet_load_path = 'path/to/your/unet_model.pth'
unet_loaded = UNet2DConditionModel(**model_config) # 创建一个与原模型结构相同的实例
unet_loaded.load_state_dict(torch.load(unet_load_path))
unet_loaded.to(device)
unet_loaded.eval() # 设置为评估模式
# 恢复优化器的状态以继续训练
opt_load_path = 'path/to/your/optimizer.pth'
opt_loaded = torch.optim.AdamW(unet_loaded.parameters(), lr=1e-4) # 创建一个新的优化器实例
opt_loaded.load_state_dict(torch.load(opt_load_path))
# 使用unet_loaded进行推理或者用opt_loaded继续训练。

复制代码

betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2

复制代码

import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
# 组合图像，生成grid
def image_grid(imgs, rows, cols):
assert len(imgs) == rows*cols
w, h = imgs[0].size
grid = Image.new('RGB', size=(cols*w, rows*h))
grid_w, grid_h = grid.size
for i, img in enumerate(imgs):
grid.paste(img, box=(i%cols*w, i//cols*h))
return grid
# 加载文生图pipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", # 或者使用 SD v1.4: "CompVis/stable-diffusion-v1-4"
torch_dtype=torch.float16
).to("cuda")
# 输入text，这里text又称为prompt
prompts = [
"a photograph of an astronaut riding a horse",
"A cute otter in a rainbow whirlpool holding shells, watercolor",
"An avocado armchair",
"A white dog wearing sunglasses"
]
generator = torch.Generator("cuda").manual_seed(42) # 定义随机seed，保证可重复性
# 执行推理
images = pipe(
prompts,
height=512,
width=512,
num_inference_steps=50,
guidance_scale=7.5,
negative_prompt=None,
num_images_per_prompt=1,
generator=generator
).images
# 保存每个单独的图片
for idx, img in enumerate(images):
img.save(f"image_{idx}.png")
# 创建并保存组合后的网格图
grid = image_grid(images, rows=1, cols=len(prompts))
grid.save("combined_images.png")
print("所有图片已保存到本地。")

复制代码

from diffusers import DDIMScheduler
# 注意这里的clip_sample要关闭，否则生成图像存在问题，因为不能对latent进行clip
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, clip_sample=False)

复制代码

import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
# 加载图生图pipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
# 读取初始图片
init_image = Image.open("liuyifei.jpg").convert("RGB").resize((512, 512))
print(init_image.size)
init_image.save("liuyifei_512.jpg")
# 推理
prompt = "A girl wearing a hat on her head."
generator = torch.Generator(device="cuda").manual_seed(2023)
image = pipe(
prompt=prompt,
image=init_image,
strength=0.8,
guidance_scale=7.5,
generator=generator
).images[0]
# 保存生成的图像
output_path = "generated_liuyifei.jpg"
image.save(output_path)
print(f"Generated image saved to {output_path}")

复制代码

import torch
from diffusers import StableDiffusionInpaintPipelineLegacy
from PIL import Image
# 加载inpainting pipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionInpaintPipelineLegacy.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
# 读取输入图像和输入mask
input_image = Image.open("overture-creations-5sI6fQgYIuo.png").resize((512, 512))
input_mask = Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512, 512))
# 执行推理
prompt = ["a mecha robot sitting on a bench", "a cat sitting on a bench"]
generator = torch.Generator("cuda").manual_seed(0)
with torch.autocast("cuda"):
images = pipe(
prompt=prompt,
image=input_image,
mask_image=input_mask,
num_inference_steps=50,
strength=0.75,
guidance_scale=7.5,
num_images_per_prompt=1,
generator=generator,
).images
# 保存每个单独的图片
for idx, img in enumerate(images):
img.save(f"image_{idx}.png")
print("所有图片已保存到本地。")

复制代码