text: import warnings
  import argparse
  import os
  from PIL import Image
  import numpy as np
  import torch
  
  import stylegan2
  from stylegan2 import utils
  
  # ----------------------------------------------------------------------------
  
  _description = """StyleGAN2 generator.
  Run 'python %(prog)s <subcommand> --help' for subcommand help."""
  
  # ----------------------------------------------------------------------------
  
  _examples = """examples:
    # Train a network or convert a pretrained one.
    # Example of converting pretrained ffhq model:
    python run_convert_from_tf --download ffhq-config-f --output G.pth D.pth Gs.pth
  
    # Generate ffhq uncurated images (matches paper Figure 12)
    python %(prog)s generate_images --network=Gs.pth --seeds=6600-6625 --truncation_psi=0.5
  
    # Generate ffhq curated images (matches paper Figure 11)
    python %(prog)s generate_images --network=Gs.pth --seeds=66,230,389,1518 --truncation_psi=1.0
  
    # Example of converting pretrained car model:
    python run_convert_from_tf --download car-config-f --output G_car.pth D_car.pth Gs_car.pth
  
    # Generate uncurated car images (matches paper Figure 12)
    python %(prog)s generate_images --network=Gs_car.pth --seeds=6000-6025 --truncation_psi=0.5
  
    # Generate style mixing example (matches style mixing video clip)
    python %(prog)s style_mixing_example --network=Gs.pth --row_seeds=85,100,75,458,1500 --col_seeds=55,821,1789,293 --truncation_psi=1.0
  """
  
  
  # ----------------------------------------------------------------------------
  
  def _add_shared_arguments(parser):
      parser.add_argument(
          '--network',
          help='Network file path',
          required=True,
          metavar='FILE'
      )
  
      parser.add_argument(
          '--output',
          help='Root directory for run results. Default: %(default)s',
          type=str,
          default='./results',
          metavar='DIR'
      )
  
      parser.add_argument(
          '--pixel_min',
          help='Minumum of the value range of pixels in generated images. ' + \
               'Default: %(default)s',
          default=-1,
          type=float,
          metavar='VALUE'
      )
  
      parser.add_argument(
          '--pixel_max',
          help='Maximum of the value range of pixels in generated images. ' + \
               'Default: %(default)s',
          default=1,
          type=float,
          metavar='VALUE'
      )
  
      parser.add_argument(
          '--gpu',
          help='CUDA device indices (given as separate ' + \
               'values if multiple, i.e. "--gpu 0 1"). Default: Use CPU',
          type=int,
          default=[],
          nargs='*',
          metavar='INDEX'
      )
  
      parser.add_argument(
          '--truncation_psi',
          help='Truncation psi. Default: %(default)s',
          type=float,
          default=0.5,
          metavar='VALUE'
      )
  
  
  def get_arg_parser():
      parser = argparse.ArgumentParser(
          description=_description,
          epilog=_examples,
          formatter_class=argparse.RawDescriptionHelpFormatter
      )
  
      range_desc = 'NOTE: This is a single argument, where list ' + \
                   'elements are separated by "," and ranges are defined as "a-b". Only integers are allowed.'
  
      subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
  
      generate_images_parser = subparsers.add_parser(
          'generate_images', help='Generate images')
  
      generate_images_parser.add_argument(
          '--batch_size',
          help='Batch size for generator. Default: %(default)s',
          type=int,
          default=1,
          metavar='VALUE'
      )
  
      generate_images_parser.add_argument(
          '--seeds',
          help='List of random seeds for generating images. ' + range_desc,
          type=utils.range_type,
          required=True,
          metavar='RANGE'
      )
  
      _add_shared_arguments(generate_images_parser)
  
      style_mixing_example_parser = subparsers.add_parser(
          'style_mixing_example', help='Generate style mixing video')
  
      style_mixing_example_parser.add_argument(
          '--row_seeds',
          help='List of random seeds for image rows. ' + range_desc,
          type=utils.range_type,
          required=True,
          metavar='RANGE'
      )
  
      style_mixing_example_parser.add_argument(
          '--col_seeds',
          help='List of random seeds for image columns. ' + range_desc,
          type=utils.range_type,
          required=True,
          metavar='RANGE'
      )
  
      style_mixing_example_parser.add_argument(
          '--style_layers',
          help='Indices of layers to mix style for. ' + \
               'Default: %(default)s. ' + range_desc,
          type=utils.range_type,
          default='0-6',
          metavar='RANGE'
      )
  
      style_mixing_example_parser.add_argument(
          '--grid',
          help='Save a grid as well of the style mix. Default: %(default)s',
          type=utils.bool_type,
          default=True,
          const=True,
          nargs='?',
          metavar='BOOL'
      )
  
      _add_shared_arguments(style_mixing_example_parser)
  
      return parser
  
  
  # ----------------------------------------------------------------------------
  
  def style_mixing_example(G, args):
      assert max(args.style_layers) < len(G), \
          'Style layer indices can not be larger than ' + \
          'number of style layers ({}) of the generator.'.format(len(G))
      device = torch.device(args.gpu[0] if args.gpu else 'cpu')
      if device.index is not None:
          torch.cuda.set_device(device.index)
      if len(args.gpu) > 1:
          warnings.warn('Multi GPU is not available for style mixing example. Using device {}'.format(device))
      G.to(device)
      G.static_noise()
      latent_size, label_size = G.latent_size, G.label_size
      G_mapping, G_synthesis = G.G_mapping, G.G_synthesis
  
      all_seeds = list(set(args.row_seeds + args.col_seeds))
      all_z = torch.stack([torch.from_numpy(np.random.RandomState(seed).randn(latent_size)) for seed in all_seeds])
      all_z = all_z.to(device=device, dtype=torch.float32)
      if label_size:
          labels = torch.zeros(len(all_z), dtype=torch.int64, device=device)
      else:
          labels = None
  
      print('Generating disentangled latents...')
      with torch.no_grad():
          all_w = G_mapping(latents=all_z, labels=labels)
  
      all_w = all_w.unsqueeze(1).repeat(1, len(G_synthesis), 1)
  
      w_avg = G.dlatent_avg
  
      if args.truncation_psi != 1:
          all_w = w_avg + args.truncation_psi * (all_w - w_avg)
  
      w_dict = {seed: w for seed, w in zip(all_seeds, all_w)}
  
      all_images = []
  
      progress = utils.ProgressWriter(len(all_w))
      progress.write('Generating images...', step=False)
  
      with torch.no_grad():
          for w in all_w:
              all_images.append(G_synthesis(w.unsqueeze(0)))
              progress.step()
  
      progress.write('Done!', step=False)
      progress.close()
  
      all_images = torch.cat(all_images, dim=0)
  
      image_dict = {(seed, seed): image for seed, image in zip(all_seeds, all_images)}
  
      progress = utils.ProgressWriter(len(args.row_seeds) * len(args.col_seeds))
      progress.write('Generating style-mixed images...', step=False)
  
      for row_seed in args.row_seeds:
          for col_seed in args.col_seeds:
              w = w_dict[row_seed].clone()
              w[args.style_layers] = w_dict[col_seed][args.style_layers]
              with torch.no_grad():
                  image_dict[(row_seed, col_seed)] = G_synthesis(w.unsqueeze(0)).squeeze(0)
              progress.step()
  
      progress.write('Done!', step=False)
      progress.close()
  
      progress = utils.ProgressWriter(len(image_dict))
      progress.write('Saving images...', step=False)
  
      for (row_seed, col_seed), image in list(image_dict.items()):
          image = utils.tensor_to_PIL(
              image, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
          image_dict[(row_seed, col_seed)] = image
          image.save(os.path.join(args.output, '%d-%d.png' % (row_seed, col_seed)))
          progress.step()
  
      progress.write('Done!', step=False)
      progress.close()
  
      if args.grid:
          print('\n\nSaving style-mixed grid...')
          H, W = all_images.size()[2:]
          canvas = Image.new(
              'RGB', (W * (len(args.col_seeds) + 1), H * (len(args.row_seeds) + 1)), 'black')
          for row_idx, row_seed in enumerate([None] + args.row_seeds):
              for col_idx, col_seed in enumerate([None] + args.col_seeds):
                  if row_seed is None and col_seed is None:
                      continue
                  key = (row_seed, col_seed)
                  if row_seed is None:
                      key = (col_seed, col_seed)
                  if col_seed is None:
                      key = (row_seed, row_seed)
                  canvas.paste(image_dict[key], (W * col_idx, H * row_idx))
          canvas.save(os.path.join(args.output, 'grid.png'))
          print('Done!')
  
  
  # ----------------------------------------------------------------------------
  
  def generate_images(G, args):
      latent_size, label_size = G.latent_size, G.label_size
      device = torch.device(args.gpu[0] if args.gpu else 'cpu')
      if device.index is not None:
          torch.cuda.set_device(device.index)
      G.to(device)
      if args.truncation_psi != 1:
          G.set_truncation(truncation_psi=args.truncation_psi)
      if len(args.gpu) > 1:
          warnings.warn(
              'Noise can not be randomized based on the seed ' + \
              'when using more than 1 GPU device. Noise will ' + \
              'now be randomized from default random state.'
          )
          G.random_noise()
          G = torch.nn.DataParallel(G, device_ids=args.gpu)
      else:
          noise_reference = G.static_noise()
  
      def get_batch(seeds):
          latents = []
          labels = []
          if len(args.gpu) <= 1:
              noise_tensors = [[] for _ in noise_reference]
          for seed in seeds:
              rnd = np.random.RandomState(seed)
              latents.append(torch.from_numpy(rnd.randn(latent_size)))
              if len(args.gpu) <= 1:
                  for i, ref in enumerate(noise_reference):
                      noise_tensors[i].append(torch.from_numpy(rnd.randn(*ref.size()[1:])))
              if label_size:
                  labels.append(torch.tensor([rnd.randint(0, label_size)]))
          latents = torch.stack(latents, dim=0).to(device=device, dtype=torch.float32)
          if labels:
              labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
          else:
              labels = None
          if len(args.gpu) <= 1:
              noise_tensors = [
                  torch.stack(noise, dim=0).to(device=device, dtype=torch.float32)
                  for noise in noise_tensors
              ]
          else:
              noise_tensors = None
          return latents, labels, noise_tensors
  
      progress = utils.ProgressWriter(len(args.seeds))
      progress.write('Generating images...', step=False)
  
      for i in range(0, len(args.seeds), args.batch_size):
          latents, labels, noise_tensors = get_batch(args.seeds[i: i + args.batch_size])
          if noise_tensors is not None:
              G.static_noise(noise_tensors=noise_tensors)
          with torch.no_grad():
              generated = G(latents, labels=labels)
          images = utils.tensor_to_PIL(
              generated, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
          for seed, img in zip(args.seeds[i: i + args.batch_size], images):
              img.save(os.path.join(args.output, 'seed%04d.png' % seed))
              progress.step()
  
      progress.write('Done!', step=False)
      progress.close()
  
  
  # ----------------------------------------------------------------------------
  
  def main():
      args = get_arg_parser().parse_args()
      assert args.command, 'Missing subcommand.'
      assert os.path.isdir(args.output) or not os.path.splitext(args.output)[-1], \
          '--output argument should specify a directory, not a file.'
      if not os.path.exists(args.output):
          os.makedirs(args.output)
  
      G = stylegan2.models.load(args.network)
      assert isinstance(G, stylegan2.models.Generator), 'Model type has to be ' + \
                                                        'stylegan2.models.Generator. Found {}.'.format(type(G))
  
      if args.command == 'generate_images':
          generate_images(G, args)
      elif args.command == 'style_mixing_example':
          style_mixing_example(G, args)
      else:
          raise TypeError('Unkown command {}'.format(args.command))
  
  
  # ----------------------------------------------------------------------------
  
  if __name__ == '__main__':
      main()
  