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