Gradio 和 W&B 集成

简介

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

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

什么是 Wandb？

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

什么是 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。

   

   在顶部安装 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 标签，任何人都可以将 HF spaces 上的 Gradio 演示直接嵌入到他们的博客、网站、文档等中。

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

7. （可选）在您的 Gradio 应用中嵌入 W&B 图表

   也可以在 Gradio 应用中嵌入 W&B 图表。为此，您可以创建包含您的图表的 W&B 报告，并将它们嵌入到 gr.HTML 块中的 Gradio 应用中。

   报告需要是公开的，并且您需要像这样将 URL 包装在 iFrame 中：

   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 在用户友好的界面中演示模型，以便在 Web 上的任何地方分享。

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

• 点击此处在 Hugging Face 上创建一个帐户。
• 在您的用户名下添加 Gradio 演示，请参阅本课程，了解如何在 Hugging Face 上设置 Gradio 演示。
• 点击此处请求加入 wandb 组织。
• 获得批准后，将模型从您的用户名转移到 Wandb 组织