Spaces:
Runtime error
Runtime error
| # 导入需要的包 | |
| 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) | |