Powered by AppSignal & Oban Pro
Would you like to see your link here? Contact us

Style-based GAN

demo_StyleGAN2-ADA/StyleGAN2ADA.livemd

Style-based GAN

File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")
System.put_env("NNCOMPILED", "YES")

Mix.install([
  {:tfl_interp, "~> 0.1.13"},
  {:cimg, "~> 0.1.19"},
  {:kino, "~> 0.7.0"},
  {:nx, "~> 0.5.3"}
])

0.Original work

“Training Generative Adversarial Networks with Limited Data”

GitHub: StyleGAN2 with adaptive discriminator augmentation

GitHub: sugyan/gan-playground

The tflite model afhqdog.tflite is converted from pretraind model in pickle format.

Thanks a lot!!!


Implementation with TflInterp in Elixir

1.Defining the inference module: StyleGAN2ADA

defmodule StyleGAN2ADA.Mapping do
  alias TflInterp, as: NNInterp

  use NNInterp,
    model: "./model/afhqdog.mapping.tflite",
    url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.mapping.tflite",
    inputs: [f32: {1, 512}],
    outputs: [f32: {1, 16.512}]

  def dlatants() do
    for _ <- 1..512, into: "" do
      <<:rand.normal()::little-float-32>>
    end
    |> Nx.from_binary(:f32)
    |> Nx.reshape({1, 512})
    |> dlatants()
  end

  def dlatants(latants) do
    # preprocess
    input0 = Nx.to_binary(latants)

    # prediction
    output0 =
      session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)

    # postprocess
    Nx.from_binary(output0, :f32) |> Nx.reshape({16, 512})
  end
end
defmodule StyleGAN2ADA.Synthesis do
  @width 512
  @height 512

  alias TflInterp, as: NNInterp

  use NNInterp,
    model: "./model/afhqdog.synthesis.tflite",
    url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.synthesis.tflite",
    inputs: [f32: {1, 16, 512}],
    outputs: [f32: {1, 3, @height, @width}]

  def image(dlatants) do
    # preprocess
    input0 = Nx.to_binary(dlatants)

    # prediction
    output0 =
      session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)

    # postprocess
    CImg.from_binary(output0, @width, @height, 1, 3, [{:range, {-1.0, 1.0}}, :nchw])
  end
end

Launch StyleGAN2ADA.Mapping and StyleGAN2ADA.Synthesis.

# TflInterp.stop(StyleGAN2ADA.Mapping)
StyleGAN2ADA.Mapping.start_link([])
# TflInterp.stop(StyleGAN2ADA.Synthesis)
StyleGAN2ADA.Synthesis.start_link([])

Display the properties of the StyleGAN2ADA model.

TflInterp.info(StyleGAN2ADA.Mapping)
TflInterp.info(StyleGAN2ADA.Synthesis)

2.Let’s try it

Morphing between two randomly generated images.

a = StyleGAN2ADA.Mapping.dlatants()
b = StyleGAN2ADA.Mapping.dlatants()
delta = Nx.subtract(b, a) |> Nx.divide(5)

Enum.map(0..5, fn i ->
  Nx.add(a, Nx.multiply(delta, i))
  |> StyleGAN2ADA.Synthesis.image()
  |> CImg.display_kino(:jpeg)
end)
|> Kino.Layout.grid(columns: 6)