9dab8ae56ce90f1b9ba6808a6e0673ab6f13103e,GANs/stylegan2/generate.py,,main,#,98

Before Change

        print(f"using style noise seed {args.seed_mix} for layers {args.mix_after}-{num_layers}.")
        rnd = np.random.RandomState(args.seed_mix)
        z2 = rnd.randn(batch_size, 512)
        style_noises = [nn.Variable((batch_size, 512)).apply(d=z)
                        for _ in range(args.mix_after)]
        style_noises += [nn.Variable((batch_size, 512)).apply(d=z2)
                         for _ in range(num_layers - args.mix_after)]
    else:
        // no style mixing (single noise / style is used)
        print(f"using style noise seed {args.seed} for entire layers.")
        style_noise = nn.Variable((batch_size, 512)).apply(d=z)
        style_noises = [style_noise for _ in range(num_layers)]

    nn.load_parameters("styleGAN2_G_params.h5")
    rgb_output = generate(batch_size, style_noises,
                          args.stochastic_seed, args.truncation_psi)
    rgb_output.forward()

    // convert to uint8 to save an image file
    image = convert_images_to_uint8(rgb_output, drange=[-1, 1])
    if args.output_filename is None:
        if not args.mixing:
            filename = f"seed{args.seed}.png"
        else:
            filename = f"seed{args.seed}_{args.seed_mix}.png"
    else:
        filename = args.output_filename

    os.makedirs(args.output_dir, exist_ok=True)
    filepath = os.path.join(args.output_dir, filename)
    imsave(filepath, image, channel_first=True)
    print(f"Genetation completed. Saved {filepath}.")


if __name__ == "__main__":

After Change

    ctx = get_extension_context(args.context)
    nn.set_default_context(ctx)

    batch_size = args.batch_size
    num_layers = 18

    rnd = np.random.RandomState(args.seed)
    z = rnd.randn(batch_size, 512)

    print("Generation started...")
    print(f"truncation value: {args.truncation_psi}")
    print(f"seed for additional noise: {args.stochastic_seed}")

    // Inference via nn.NdArray utilizes significantly less memory

    if args.mixing:
        // apply style mixing
        assert args.seed_mix
        print(f"using style noise seed {args.seed} for layers 0-{args.mix_after - 1}")
        print(f"using style noise seed {args.seed_mix} for layers {args.mix_after}-{num_layers}.")
        rnd = np.random.RandomState(args.seed_mix)
        z2 = rnd.randn(batch_size, 512)
        style_noises = [nn.NdArray.from_numpy_array(z)]
        style_noises += [nn.NdArray.from_numpy_array(z2)]
    else:
        // no style mixing (single noise / style is used)
        print(f"using style noise seed {args.seed} for entire layers.")
        style_noises = [nn.NdArray.from_numpy_array(z) for _ in range(2)]

    nn.set_auto_forward(True)
    nn.load_parameters("styleGAN2_G_params.h5")
    rgb_output = generate(batch_size, style_noises,
                          args.stochastic_seed, args.mix_after, args.truncation_psi)

    // convert to uint8 to save an image file
    image = convert_images_to_uint8(rgb_output, drange=[-1, 1])
    if args.output_filename is None:
        if not args.mixing:
            filename = f"seed{args.seed}"
        else:
            filename = f"seed{args.seed}_{args.seed_mix}"
    else:
        filename = args.output_filename

    os.makedirs(args.output_dir, exist_ok=True)

    for i in range(batch_size):
        filepath = os.path.join(args.output_dir, f"{filename}_{i}.png")
        imsave(filepath, image[i], channel_first=True)
        print(f"Genetation completed. Saved {filepath}.")


if __name__ == "__main__":
    main()

Instances