code: import argparse
  import math
  import os
  
  import torch
  from torch import optim
  from torch.nn import functional as F
  from torchvision import transforms
  from PIL import Image
  from utils import tensor2image, save_image
  from tqdm import tqdm
  
  import lpips
  from model import Generator, Encoder
  
  
  def noise_regularize(noises):
      loss = 0
  
      for noise in noises:
          size = noise.shape[2]
  
          while True:
              loss = (
                  loss
                  + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)
                  + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)
              )
  
              if size <= 8:
                  break
  
              noise = noise.reshape([-1, 1, size // 2, 2, size // 2, 2])
              noise = noise.mean([3, 5])
              size //= 2
  
      return loss
  
  
  def noise_normalize_(noises):
      for noise in noises:
          mean = noise.mean()
          std = noise.std()
  
          noise.data.add_(-mean).div_(std)
  
  
  def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
      lr_ramp = min(1, (1 - t) / rampdown)
      lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
      lr_ramp = lr_ramp * min(1, t / rampup)
  
      return initial_lr * lr_ramp
  
  
  def latent_noise(latent, strength):
      noise = torch.randn_like(latent) * strength
  
      return latent + noise
  
  
  def make_image(tensor):
      return (
          tensor.detach()
          .clamp_(min=-1, max=1)
          .add(1)
          .div_(2)
          .mul(255)
          .type(torch.uint8)
          .permute(0, 2, 3, 1)
          .to("cpu")
          .numpy()
      )
  
  
  if __name__ == "__main__":
      device = "cuda"
  
      # -----------------------------------
      # Parser
      # -----------------------------------
  
      parser = argparse.ArgumentParser(
          description="Image projector to the generator latent spaces"
      )
      parser.add_argument(
          "--ckpt", type=str, required=True, help="path to the model checkpoint"
      )
      parser.add_argument(
          "--e_ckpt", type=str, default=None, help="path to the encoder checkpoint"
      )
      parser.add_argument(
          "--size", type=int, default=256, help="output image sizes of the generator"
      )
      parser.add_argument(
          "--truncation", type=float, default=0.7, help="truncation"
      )
      parser.add_argument(
          "--lr_rampup",
          type=float,
          default=0.05,
          help="duration of the learning rate warmup",
      )
      parser.add_argument(
          "--lr_rampdown",
          type=float,
          default=0.25,
          help="duration of the learning rate decay",
      )
      parser.add_argument("--lr", type=float, default=0.01, help="learning rate")
      parser.add_argument(
          "--noise", type=float, default=0.05, help="strength of the noise level"
      )
      parser.add_argument(
          "--noise_ramp",
          type=float,
          default=0.75,
          help="duration of the noise level decay",
      )
      parser.add_argument("--step", type=int, default=1000, help="optimize iterations")
      parser.add_argument(
          "--noise_regularize",
          type=float,
          default=1e5,
          help="weight of the noise regularization",
      )
      parser.add_argument("--mse", type=float, default=0, help="weight of the mse loss")
      parser.add_argument("--vgg", type=float, default=1.0, help="weight of the vgg loss")
      parser.add_argument(
          "--w_plus",
          action="store_true",
          help="allow to use distinct latent codes to each layers",
      )
      parser.add_argument(
          "--project_name", type=str, default="project", help="name of the result project file"
      )
      parser.add_argument(
          "--factor_name", type=str, default="factor", help="name of the result factor file"
      )
      parser.add_argument(
          "--files", nargs="+", help="path to image files to be projected"
      )
  
      args = parser.parse_args()
  
      # =============================================
  
      # -----------------------------------
      # Project Images to Latent spaces
      # -----------------------------------
      
      if args.files is None:
          exit() 
  
      n_mean_latent = 10000
  
      # Load Real Images
      resize = min(args.size, 256)
  
      transform = transforms.Compose(
          [
              transforms.Resize(resize),
              transforms.CenterCrop(resize),
              transforms.ToTensor(),
              transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
          ]
      )
  
      imgs = []
  
      for imgfile in args.files:
          img = transform(Image.open(imgfile).convert("RGB"))
          imgs.append(img)
  
      imgs = torch.stack(imgs, 0).to(device)
  
  
  
      # -------------
      # Generator
      # -------------
  
      g_ema = Generator(args.size, 512, 8).to(device)
      g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False)
      g_ema.eval()
      
      trunc = g_ema.mean_latent(4096).detach().clone()
  
      # -------------
      # Encoder
      # -------------
  
      if args.e_ckpt is not None :
          e_ckpt = torch.load(args.e_ckpt, map_location=device)
  
          encoder = Encoder(args.size, 512).to(device)
          encoder.load_state_dict(e_ckpt['e'])
          encoder.eval()
  
  
      # -------------
      # Latent vector
      # -------------
  
      if args.e_ckpt is not None :
          with torch.no_grad(): 
              latent_init = encoder(imgs)
          latent_in = latent_init.detach().clone() 
      else :
          with torch.no_grad():
              noise_sample = torch.randn(n_mean_latent, 512, device=device)
              latent_out = g_ema.style(noise_sample)
  
              latent_mean = latent_out.mean(0)
              latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5
  
          latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
  
          if args.w_plus:
              latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1)
  
      latent_in.requires_grad = True
  
      # -------------
      # Noise
      # -------------
  
      noises_single = g_ema.make_noise()
      noises = []
      for noise in noises_single:
          noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
  
      for noise in noises:
          noise.requires_grad = True
  
  
      # -------------
      # Loss
      # -------------
  
      # PerceptualLoss
      percept = lpips.PerceptualLoss(
          model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
      )
  
  
      # Optimizer
      if args.e_ckpt is not None :
          optimizer = optim.Adam([latent_in], lr=args.lr)
      else:
          optimizer = optim.Adam([latent_in] + noises, lr=args.lr)
  
      pbar = tqdm(range(args.step))
      latent_path = []
      proj_images = []
  
      # Training !
  
      for i in pbar:
  
          t = i / args.step
          lr = get_lr(t, args.lr)
  
          optimizer.param_groups[0]["lr"] = lr
  
          # fake image
          if args.e_ckpt is not None :
              img_gen, _ = g_ema([latent_in], input_is_latent=True,
                                  truncation=args.truncation, truncation_latent = trunc,
                                  randomize_noise=False)
          else:
              noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2
              latent_n = latent_noise(latent_in, noise_strength.item())
  
              img_gen, _ = g_ema([latent_n], input_is_latent=True, noise=noises)
          
          #
          batch, channel, height, width = img_gen.shape
          if height > 256:
              factor = height // 256
  
              img_gen = img_gen.reshape(
                  batch, channel, height // factor, factor, width // factor, factor
              )
              img_gen = img_gen.mean([3, 5])
          
  
          # latent
          if args.e_ckpt is not None :
              latent_hat = encoder(img_gen)
  
  
          # Loss
          p_loss = percept(img_gen, imgs).sum()        
          r_loss = torch.mean((img_gen - imgs) ** 2)       
          mse_loss = F.mse_loss(img_gen, imgs)
          
          n_loss = noise_regularize(noises)
  
          if args.e_ckpt is not None :
              style_loss = F.mse_loss(latent_hat, latent_init)
              loss = args.vgg * p_loss + r_loss + style_loss + args.mse * mse_loss
          else :
              style_loss = 0.0
              loss = args.vgg * p_loss + r_loss + args.mse * mse_loss + args.noise_regularize * n_loss 
  
  
          # update
          optimizer.zero_grad()
          loss.backward(retain_graph=True)
          optimizer.step()
  
          noise_normalize_(noises)
  
          if (i + 1) % 100 == 0:
              latent_path.append(latent_in.detach().clone())
              proj_images.append(img_gen)
  
          pbar.set_description(
              (
                  f"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f}; "
                  f"reconstruction: {r_loss:.4f}; "
                  f"mse_img: {mse_loss.item():.4f}; mse_latent: {style_loss:.4f}; lr: {lr:.4f} |"
              )
          )
  
      # =============================================
  
      # -----------------------------------
      # Save image, latent, noise
      # -----------------------------------
  
  
      # final generated image
      if args.e_ckpt is not None :
          img_gen, _ = g_ema([latent_path[-1]], input_is_latent=True,
                              truncation=args.truncation, truncation_latent = trunc,
                              randomize_noise=None)
      else:
          img_gen, _ = g_ema([latent_path[-1]], input_is_latent=True, noise=noises)
  
  
      filename = f"{args.project_name}.pt"
      img_ar = make_image(img_gen)
  
  
      images = []
      for i in range(len(proj_images)):
          img = proj_images[i][0]
          for k in range(1, len(proj_images[0])): 
              # img : torch.Size([3, 256*num_img, 256])
              img = torch.cat([img, proj_images[i][k]], dim =1) 
          images.append(img) 
  
  
      result_file = {}
      for i, input_name in enumerate(args.files):
          noise_single = []
          for noise in noises:
              noise_single.append(noise)
  
          name = os.path.splitext(os.path.basename(input_name))[0]
          result_file[name] = {
              "r_img": tensor2image(imgs[i]),
              "f_img": tensor2image(img_gen[i]),
              "p_img" : tensor2image(torch.cat(images, dim=2)),
              "latent": latent_in[i].unsqueeze(0),
              "noise": noise_single,
              "args" : args,
          }
  
          img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project.png"
          pil_img = Image.fromarray(img_ar[i])
          pil_img.save(img_name)
  
          img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project-interpolation.png"
          save_image(tensor2image(torch.cat(images, dim=2)), size = 20, out=img_name)
  
      torch.save(result_file, filename)
  
apis: torch.nn.functional.mse_loss torch.stack torch.randn_like torch.load torch.randn torch.roll torch.save torch.no_grad torch.optim.Adam torch.cat torch.mean