Spaces:
Runtime error
Runtime error
File size: 3,903 Bytes
1a93617 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# 导入需要的包
import gradio as gr
import numpy as np
import torch
import torchvision.utils
from PIL import Image, ImageColor
from tqdm import tqdm
from diffusers import DDPMPipeline, DDIMScheduler
# 检测可用的device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("*" * 10 + " device " + "*" * 10)
print(device)
# 载入一个预训练过的管线
pipeline_name = "johnowhitaker/sd-class-wikiart-from-bedrooms"
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
# 使用DDIM调度器，仅用40步生成一些图片
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
# 这里使用稍微多一些的步数
scheduler.set_timesteps(50)
# 给定一个RGB值，返回一个损失值，用于衡量图片的像素值与目标颜色相差多少：这里的目标颜色是一种浅蓝绿色，对应的RGB值为（0.1， 0.9， 0.5）
def color_loss(images, target_color=(0.1, 0.9, 0.5)):
# torch.ToTensor()取值范围是[0, 1]
# 首先对target_color进行归一化，使它的取值区间为（-1， 1）
target = (
torch.tensor(target_color).to(images.device) * 2 - 1
)
# 将所生成目标张量的形状改为（b, c, h, w）,以适配输入图像images的张量形状。
target = target[
None, :, None, None
]
# 计算图片的像素值以及目标颜色的均方误差
# abs()：求绝对值
# mean()：求平均值
error = torch.abs(
images - target
).mean()
return error
# generate(颜色， 引导损失强度) 函数用于生成图片
def generate(color, guidance_loss_scale):
# 将得到的颜色字符串 color 转换为 “RGB” 模式
target_color = ImageColor.getcolor(color, "RGB")
# 目标颜色值在[0, 1]
target_color = [a / 255 for a in target_color]
# 使用1幅 随机噪声图像 进行循环采样
x = torch.randn(4, 3, 256, 256).to(device)
# tqdm()：显示进度条 enumerate()：返回数据和数据索引下标
for i, t in tqdm(enumerate(scheduler.timesteps)):
# 对随机噪声图像添加时间步并作为模型输入
model_input = scheduler.scale_model_input(x, t)
# 预测噪声
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
# 设置输入图像的requires_grad属性为True
x = x.detach().requires_grad_()
# 模型输出当前时间步“去噪”后的图像
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# 计算损失值 * 引导损失强度
loss = color_loss(x0, target_color) * guidance_loss_scale
# 获取梯度
con_grad = -torch.autograd.grad(loss, x)[0]
# 根据梯度修改x
x = x.detach() + con_grad
# 使用调度器更新x
x = scheduler.step(noise_pred, t, x).prev_sample
# 查看结果，使用网格显示图像，每行显示4幅图像。
grid = torchvision.utils.make_grid(x, nrow=4)
# [0, 1]
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
# # np.array()：转换为数组
# # np.astype()：强转数据类型
# # Image.fromarray()：array 转化为 Image
im = Image.fromarray(np.array(im*255).astype(np.uint8))
# 保存图片test.jpeg，格式为jpeg
im.save("test.jpeg")
# 返回图片
return im
# 输入
inputs = [
# 颜色选择器
gr.ColorPicker(label="color", value="55FFAA"),
# 滑动条
gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3)
]
# 输出图像
outputs = gr.Image(label="result")
# 演示程序(demonstrate)的接口
demo = gr.Interface(
fn=generate,
inputs=inputs,
outputs=outputs,
# 示例
examples=[
["#BB2266", 3],
["#44CCAA", 5]
]
)
# 通过设置debug=True,你将能够在CoLab平台上看到错误信息
demo.launch(debug=True)
|