1. 其他教程
  2. Gradio 和 Wandb 集成

相关空间

akhaliq/JoJoGAN

Gradio 和 W&B 集成

简介

在本指南中，我们将带你了解

  • Gradio、Hugging Face Spaces 和 Wandb 简介
  • 如何使用 Wandb 集成设置 JoJoGAN 的 Gradio 演示
  • 如何在 wandb 上跟踪实验后，将你自己的 Gradio 演示贡献给 Hugging Face 上的 Wandb 组织

Wandb 是什么？

Weights and Biases (W&B) 允许数据科学家和机器学习科学家跟踪他们的机器学习实验的每个阶段，从训练到生产。任何指标都可以通过样本聚合，并以面板的形式显示在可自定义且可搜索的仪表板中，如下所示

Screen Shot 2022-08-01 at 5 54 59 PM

什么是 Hugging Face Spaces 和 Gradio？

Gradio

Gradio 允许用户通过几行 Python 代码将他们的机器学习模型演示为 Web 应用程序。Gradio 将任何 Python 函数（例如机器学习模型的推理函数）包装到用户界面中，并且可以在 jupyter notebook、colab notebook 中启动演示，也可以嵌入到你自己的网站中，并免费托管在 Hugging Face Spaces 上。

点击这里开始

Hugging Face Spaces

Hugging Face Spaces 是 Gradio 演示的免费托管选项。Spaces 提供 3 种 SDK 选项：Gradio、Streamlit 和静态 HTML 演示。Spaces 可以是公共的或私有的，工作流程类似于 github 存储库。Hugging Face 上目前有 2000 多个 spaces。点击这里了解更多关于 Spaces 的信息。

设置 JoJoGAN 的 Gradio 演示

现在，让我们来指导你如何自己完成这项工作。在本教程中，我们将假设你是 W&B 和 Gradio 的新手。

让我们开始吧！

  1. 创建 W&B 账户

    如果你还没有账户，请按照这些快速说明创建免费账户。这应该只需要几分钟。完成后（或者如果你已经有账户了），接下来我们将运行一个快速 colab。

  2. 打开 Colab 安装 Gradio 和 W&B

    我们将遵循 JoJoGAN 存储库中提供的 colab，并进行一些小的修改，以更有效地使用 Wandb 和 Gradio。

    Open In Colab

    在顶部安装 Gradio 和 Wandb

    pip install gradio wandb

  3. 微调 StyleGAN 和 W&B 实验跟踪

    下一步将打开一个 W&B 仪表板来跟踪你的实验，以及一个 Gradio 面板，其中包含一个下拉菜单，可从托管在 Huggingface Spaces 上的 Gradio 演示中选择预训练模型。这是你需要的代码

    alpha =  1.0
alpha = 1-alpha

preserve_color = True
num_iter = 100
log_interval = 50

samples = []
column_names = ["Reference (y)", "Style Code(w)", "Real Face Image(x)"]

wandb.init(project="JoJoGAN")
config = wandb.config
config.num_iter = num_iter
config.preserve_color = preserve_color
wandb.log(
{"Style reference": [wandb.Image(transforms.ToPILImage()(target_im))]},
step=0)

# load discriminator for perceptual loss
discriminator = Discriminator(1024, 2).eval().to(device)
ckpt = torch.load('models/stylegan2-ffhq-config-f.pt', map_location=lambda storage, loc: storage)
discriminator.load_state_dict(ckpt["d"], strict=False)

# reset generator
del generator
generator = deepcopy(original_generator)

g_optim = optim.Adam(generator.parameters(), lr=2e-3, betas=(0, 0.99))

# Which layers to swap for generating a family of plausible real images -> fake image
if preserve_color:
    id_swap = [9,11,15,16,17]
else:
    id_swap = list(range(7, generator.n_latent))

for idx in tqdm(range(num_iter)):
    mean_w = generator.get_latent(torch.randn([latents.size(0), latent_dim]).to(device)).unsqueeze(1).repeat(1, generator.n_latent, 1)
    in_latent = latents.clone()
    in_latent[:, id_swap] = alpha*latents[:, id_swap] + (1-alpha)*mean_w[:, id_swap]

    img = generator(in_latent, input_is_latent=True)

    with torch.no_grad():
        real_feat = discriminator(targets)
    fake_feat = discriminator(img)

    loss = sum([F.l1_loss(a, b) for a, b in zip(fake_feat, real_feat)])/len(fake_feat)

    wandb.log({"loss": loss}, step=idx)
    if idx % log_interval == 0:
        generator.eval()
        my_sample = generator(my_w, input_is_latent=True)
        generator.train()
        my_sample = transforms.ToPILImage()(utils.make_grid(my_sample, normalize=True, range=(-1, 1)))
        wandb.log(
        {"Current stylization": [wandb.Image(my_sample)]},
        step=idx)
    table_data = [
            wandb.Image(transforms.ToPILImage()(target_im)),
            wandb.Image(img),
            wandb.Image(my_sample),
        ]
    samples.append(table_data)

    g_optim.zero_grad()
    loss.backward()
    g_optim.step()

out_table = wandb.Table(data=samples, columns=column_names)
wandb.log({"Current Samples": out_table})

  4. 保存、下载和加载模型

    以下是保存和下载模型的方法。

    from PIL import Image
import torch
torch.backends.cudnn.benchmark = True
from torchvision import transforms, utils
from util import *
import math
import random
import numpy as np
from torch import nn, autograd, optim
from torch.nn import functional as F
from tqdm import tqdm
import lpips
from model import *
from e4e_projection import projection as e4e_projection

from copy import deepcopy
import imageio

import os
import sys
import torchvision.transforms as transforms
from argparse import Namespace
from e4e.models.psp import pSp
from util import *
from huggingface_hub import hf_hub_download
from google.colab import files

torch.save({"g": generator.state_dict()}, "your-model-name.pt")

files.download('your-model-name.pt')

latent_dim = 512
device="cuda"
model_path_s = hf_hub_download(repo_id="akhaliq/jojogan-stylegan2-ffhq-config-f", filename="stylegan2-ffhq-config-f.pt")
original_generator = Generator(1024, latent_dim, 8, 2).to(device)
ckpt = torch.load(model_path_s, map_location=lambda storage, loc: storage)
original_generator.load_state_dict(ckpt["g_ema"], strict=False)
mean_latent = original_generator.mean_latent(10000)

generator = deepcopy(original_generator)

ckpt = torch.load("/content/JoJoGAN/your-model-name.pt", map_location=lambda storage, loc: storage)
generator.load_state_dict(ckpt["g"], strict=False)
generator.eval()

plt.rcParams['figure.dpi'] = 150

transform = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

def inference(img):
    img.save('out.jpg')
    aligned_face = align_face('out.jpg')

    my_w = e4e_projection(aligned_face, "out.pt", device).unsqueeze(0)
    with torch.no_grad():
        my_sample = generator(my_w, input_is_latent=True)

    npimage = my_sample[0].cpu().permute(1, 2, 0).detach().numpy()
    imageio.imwrite('filename.jpeg', npimage)
    return 'filename.jpeg'

  5. 构建 Gradio 演示

    import gradio as gr

title = "JoJoGAN"
description = "Gradio Demo for JoJoGAN: One Shot Face Stylization. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."

demo = gr.Interface(
    inference,
    gr.Image(type="pil"),
    gr.Image(type="file"),
    title=title,
    description=description
)

demo.launch(share=True)

  6. 将 Gradio 演示集成到你的 W&B 仪表板中

    最后一步——将 Gradio 演示与 W&B 仪表板集成——只需多一行代码

    demo.integrate(wandb=wandb)

    一旦你调用 integrate，就会创建一个演示，你可以将其集成到你的仪表板或报告中。

    在 W&B 之外使用 Web 组件，使用 gradio-app 标签，任何人都可以将 Hugging Face Spaces 上的 Gradio 演示直接嵌入到他们的博客、网站、文档等中。

    <gradio-app space="akhaliq/JoJoGAN"> </gradio-app>

  7. （可选）在你的 Gradio App 中嵌入 W&B 图表

    也可以在 Gradio 应用程序中嵌入 W&B 图表。为此，你可以创建图表的 W&B 报告，并将其嵌入 Gradio 应用程序中的 gr.HTML 块中。

    报告需要是公开的，并且你需要使用 iFrame 包装 URL，如下所示

    import gradio as gr

def wandb_report(url):
    iframe = f'<iframe src={url} style="border:none;height:1024px;width:100%">'
    return gr.HTML(iframe)

with gr.Blocks() as demo:
    report_url = 'https://wandb.ai/_scott/pytorch-sweeps-demo/reports/loss-22-10-07-16-00-17---VmlldzoyNzU2NzAx'
    report = wandb_report(report_url)

demo.launch(share=True)

结论

我们希望你喜欢这个将 Gradio 演示嵌入 W&B 报告的简短演示！感谢你坚持到最后。总结一下

  • 微调 JoJoGAN 只需要一张参考图像，通常在 colab 中的 GPU 上需要大约 1 分钟。训练完成后，可以将样式应用于任何输入图像。在论文中阅读更多内容。

  • W&B 只需向 colab 添加几行代码即可跟踪实验，你可以在一个集中的仪表板中可视化、排序和理解你的实验。

  • 同时，Gradio 在用户友好的界面中演示模型，以便在网络上的任何地方共享。

如何在 Wandb 组织下的 HF spaces 上贡献 Gradio 演示

  • 点击这里在 Hugging Face 上创建账户。
  • 在你的用户名下添加 Gradio 演示，请参阅本课程以了解如何在 Hugging Face 上设置 Gradio 演示。
  • 点击这里请求加入 wandb 组织。
  • 批准后，将模型从你的用户名转移到 Wandb 组织
←

Gradio 和 Hugging Face 上的 ONNX

创建沉浸式演示

→
gradio