zwbdla commited on
Commit
0f38d22
·
1 Parent(s): d84e317

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -110
app.py DELETED
@@ -1,110 +0,0 @@
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)