Gradio 和 W&B 集成

介绍

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

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

什么是Wandb？

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

什么是Hugging Face Spaces和Gradio？

Gradio

Gradio 允许用户将他们的机器学习模型作为网页应用进行演示，只需几行 Python 代码。Gradio 将任何 Python 函数（例如机器学习模型的推理函数）包装成用户界面，并且这些演示可以在 Jupyter 笔记本、Colab 笔记本中启动，也可以嵌入到您自己的网站中，并免费托管在 Hugging Face Spaces 上。

从这里开始 here

Hugging Face 空间

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. 要使用它，只需上传您的图像，或点击其中一个示例来加载它们。在下面的链接中阅读更多信息。"
   
   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空间中的Gradio演示直接嵌入到他们的博客、网站、文档等中：

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

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

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

   报告需要公开，并且你需要将URL包装在一个iFrame中，如下所示：

   import gradio as gr
   
   def wandb_report(url):
       iframe = f'
       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空间上贡献Gradio演示

• 在Hugging Face上创建一个账户 这里.
• 在你的用户名下添加Gradio演示，请参阅此课程以在Hugging Face上设置Gradio演示。
• 请求加入wandb组织 这里.
• 一旦批准，将模型从您的用户名转移到Wandb组织