<jupyter_start><jupyter_text>Fetch Codebase<jupyter_code>#@title
import os
os.chdir('/content')
CODE_DIR = 'sefa'
!git clone https://github.com/genforce/sefa.git $CODE_DIR
os.chdir(f'./{CODE_DIR}')<jupyter_output><empty_output><jupyter_text>Define Utility Functions<jupyter_code>#@title
import os.path
import io
import IPython.display
import numpy as np
import cv2
import PIL.Image

import torch

from models import parse_gan_type
from utils import to_tensor
from utils import postprocess
from utils import load_generator
from utils import factorize_weight


def sample(generator, gan_type, num=1, seed=0):
  """Samples latent codes."""
  torch.manual_seed(seed)
  codes = torch.randn(num, generator.z_space_dim).cuda()
  if gan_type == 'pggan':
    codes = generator.layer0.pixel_norm(codes)
  elif gan_type == 'stylegan':
    codes = generator.mapping(codes)['w']
    codes = generator.truncation(codes, trunc_psi=0.7, trunc_layers=8)
  elif gan_type == 'stylegan2':
    codes = generator.mapping(codes)['w']
    codes = generator.truncation(codes, trunc_psi=0.5, trunc_layers=18)
  codes = codes.detach().cpu().numpy()
  return codes


def synthesize(generator, gan_type, codes):
  """Synthesizes images with the give codes."""
  if gan_type == 'pggan':
    images = generator(to_tensor(codes))['image']
  elif gan_type in ['stylegan', 'stylegan2']:
    images = generator.synthesis(to_tensor(codes))['image']
  images = postprocess(images)
  return images


def imshow(images, col, viz_size=256):
  """Shows images in one figure."""
  num, height, width, channels = images.shape
  assert num % col == 0
  row = num // col

  fused_image = np.zeros((viz_size * row, viz_size * col, channels), dtype=np.uint8)

  for idx, image in enumerate(images):
    i, j = divmod(idx, col)
    y = i * viz_size
    x = j * viz_size
    if height != viz_size or width != viz_size:
      image = cv2.resize(image, (viz_size, viz_size))
    fused_image[y:y + viz_size, x:x + viz_size] = image

  fused_image = np.asarray(fused_image, dtype=np.uint8)
  data = io.BytesIO()
  PIL.Image.fromarray(fused_image).save(data, 'jpeg')
  im_data = data.getvalue()
  disp = IPython.display.display(IPython.display.Image(im_data))
  return disp<jupyter_output><empty_output><jupyter_text>Select a Model<jupyter_code>#@title { display-mode: "form", run: "auto" }
model_name = "stylegan_animeface512" #@param ['stylegan_animeface512', 'stylegan_car512', 'stylegan_cat256', 'pggan_celebahq1024', 'stylegan_bedroom256']

generator = load_generator(model_name)
gan_type = parse_gan_type(generator)<jupyter_output><empty_output><jupyter_text>Sample Latent Codes<jupyter_code>#@title { display-mode: "form", run: "auto" }

num_samples = 3 #@param {type:"slider", min:1, max:8, step:1}
noise_seed = 0 #@param {type:"slider", min:0, max:1000, step:1}

codes = sample(generator, gan_type, num_samples, noise_seed)
images = synthesize(generator, gan_type, codes)
imshow(images, col=num_samples)<jupyter_output><empty_output><jupyter_text>Factorize & Edit<jupyter_code>#@title { display-mode: "form", run: "auto" }

layer_idx = "0-1" #@param ['all', '0-1', '2-5', '6-13']
semantic_1 = 0 #@param {type:"slider", min:-3.0, max:3.0, step:0.1}
semantic_2 = 0 #@param {type:"slider", min:-3.0, max:3.0, step:0.1}
semantic_3 = 0 #@param {type:"slider", min:-3.0, max:3.0, step:0.1}
semantic_4 = 0 #@param {type:"slider", min:-3.0, max:3.0, step:0.1}
semantic_5 = 0 #@param {type:"slider", min:-3.0, max:3.0, step:0.1}

# Fast implementation to factorize the weight by SeFa.
layers, boundaries, _ = factorize_weight(generator, layer_idx)

new_codes = codes.copy()
for sem_idx in range(5):
  boundary = boundaries[sem_idx:sem_idx + 1]
  step = eval(f'semantic_{sem_idx + 1}')
  if gan_type == 'pggan':
    new_codes += boundary * step
  elif gan_type in ['stylegan', 'stylegan2']:
    new_codes[:, layers, :] += boundary * step
new_images = synthesize(generator, gan_type, new_codes)
imshow(new_images, col=num_samples)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14985014985014986, 'fraction_numerical': 0.025724275724275724, 'mean_word_length': 5.501623376623376, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '407010', 'license': 'MIT', 'n_tokens_mistral': 1544, 'n_tokens_neox': 1480, 'n_words': 327, 'path': 'docs/SeFa.ipynb', 'repo_name': 'NoirChad/sefa', 'size': '4004'}