[Code Review] Cartoon-StyleGAN: Fine-tuning StyleGAN2 for Cartoon Face Generation

cartoon_stylegan2.ipynb 코드 리뷰

https://github.com/happy-jihye/Cartoon-StyleGAN/blob/main/Cartoon_StyleGAN2.ipynb

GitHub - happy-jihye/Cartoon-StyleGAN: Fine-tuning StyleGAN2 for Cartoon Face Generation

 

1. Set up and Generate Images using pretrained model

• Imges 생성 환경 : cuda gpu사용, 샘플 5개 생성, truncation(샘플 다양성, 품질 간의 trade-off 결정변수) 0.7
• 사전훈련된 4개의 generator 로드

import torch
import cv2 
import numpy as np
from torchvision.utils import save_image, make_grid
import PIL.Image as pilimg

from model import Generator


device='cuda'
n_sample=5
truncation = 0.7


# =============================================

# Genearaotr1
network1 = 'ffhq256'  #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network1 = f'/content/Cartoon-StyleGan2/networks/{network1}.pt' 
network1 = torch.load(network1)

g1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g1.load_state_dict(network1["g_ema"], strict=False)
trunc1 = g1.mean_latent(4096)

# Generator2
network2 = 'NaverWebtoon_StructureLoss' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network2 = f'/content/Cartoon-StyleGan2/networks/{network2}.pt' 
network2 = torch.load(network2)

g2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g2.load_state_dict(network2["g_ema"], strict=False)
trunc2 = g2.mean_latent(4096)

# Generator3
network3 = 'Disney_StructureLoss' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network3 = f'/content/Cartoon-StyleGan2/networks/{network3}.pt' 
network3 = torch.load(network3)

g3 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g3.load_state_dict(network3["g_ema"], strict=False)
trunc3 = g3.mean_latent(4096)

# Generator4
network4 = 'Metface_StructureLoss' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network4 = f'/content/Cartoon-StyleGan2/networks/{network4}.pt' 
network4 = torch.load(network4)

g4 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g4.load_state_dict(network4["g_ema"], strict=False)
trunc4 = g4.mean_latent(4096)

 

 

# 이미지 생성 후 저장

• 생성할 이미지수 11, 이미지 사이의 보간 단계 수 7
• 생성기에서 레이어를 교체하기 위한  옵션설정 (type=slider 유형은 뭐지?)
• 임의의 latent1, latent2 잠재 벡터를 생성하고 벡터 사이를 보간하는데 사용되는 latent_interp 변수 생성
• 단계마다 보간된 잠재벡터(latent_interp) 와 추가 옵션으로 g1, g2, g3, g4로 이미지 생성
• 생성된 이미지는 그리드로 연결되어 디스크에 저장, latent2벡터를 다음 시작점으로 하여 나머지 이미지도 동일하게 반복

# directory to save image
outdir = 'results' #@param {type:"string"}
if not os.path.isdir(f'{outdir}'):
   os.makedirs(f'./asset/{outdir}', exist_ok=True)

imgs = []
number_of_img = 11 #@param {type:"slider", min:0, max:30, step:1}
number_of_step = 7 #@param {type:"slider", min:0, max:10, step:1}

swap = True #@param {type:"boolean"}
swap_layer_num = 2 #@param {type:"slider", min:1, max:6, step:1}

with torch.no_grad():

  latent1 = torch.randn(1, 14, 512, device=device)
  latent1 = g1.get_latent(latent1)
  latent_interp = torch.zeros(1, latent1.shape[1], latent1.shape[2]).to(device)

  for i in range(number_of_img):
      # latent1

      latent2 = torch.randn(1, 14, 512, device=device)
      latent2 = g1.get_latent(latent2)


      for j in range(number_of_step):

          latent_interp = latent1 + (latent2-latent1) * float(j/(number_of_step-1))

          imgs_gen1, save_swap_layer = g1([latent_interp],
                                  input_is_latent=True,                                     
                                  truncation=0.7,
                                  truncation_latent=trunc1,
                                  swap=swap, swap_layer_num=swap_layer_num,
                                  randomize_noise=False)
          imgs_gen2, _ = g2([latent_interp],
       						input_is_latent=True,                                     
                            truncation=0.7,
                           	truncation_latent=trunc2,
                            swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
                                  )
          imgs_gen3, _ = g3([latent_interp],
                                  input_is_latent=True,                                     
                                  truncation=0.7,
                                  truncation_latent=trunc3,
                                  swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
                                  )
          imgs_gen4, _ = g4([latent_interp],
                                  input_is_latent=True,                                     
                                  truncation=0.7,
                                  truncation_latent=trunc4,
                                  swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
                                  )

                          
          grid = make_grid(torch.cat([imgs_gen1, imgs_gen2, imgs_gen3, imgs_gen4], 0),
                            nrow=4,
                            normalize=True,
                            range=(-1,1),
                            )
          
          ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
          im = pilimg.fromarray(ndarr)
          im.save(f'/content/Cartoon-StyleGan2/asset/{outdir}/out-{i*number_of_step+j}.png')


    latent1 = latent2

의문점 1) latent1, latent2 벡터 사이를 보간하는 latent_interp 벡터로 이미지를 생성하는 이유가 뭘까?

-> 보간된 잠재 벡터는 무작위로 생성되는 두 개의 다른 잠재 벡터 사이에 있는 벡터로 두 벡터와 공간은 동일하다. 보간된 벡터를 사용하면 두 원본 이미지 의미있는 속성을 갖는 새 이미지를 생성할 수 있다. 이 코드에서는 latent1, latent2의 가중평균을 취하여 보간 잠재 벡터를 생성하게 된다. 

의문점 2) 왜 latent2를 다음 이미지 생성할 때 사용하는걸까?

의문점 3) 두 개의 이미지로부터 합성된 이미지를 생성하는 것인데 왜 generator가 4개가 쓰이는 걸까? 2개만 있어도 되는 거 아닌가?

 

 

# 생성된 이미지를 비디오로 저장

• 첫 번째 이미지의 프레임 크기를 가져온 후 초당 프레임 수와 비디오 코덱 설정
• 생성된 모든 이미지에 대해서 반복하여 비디오에 추가
• HTML 코드를 사용하여 colab에 동영상 표시 

# Setting for Video

from PIL import Image

im = Image.open(f'/content/Cartoon-StyleGan2/asset/{outdir}/out-0.png')
frameSize = im.size
FPS = 11 #@param {type:"slider", min:0, max:30, step:1}
fourcc = cv2.VideoWriter_fourcc(*'XVID')
video_name = "results8" #@param {type:"string"}

video = cv2.VideoWriter(f'/content/Cartoon-StyleGan2/asset/{outdir}/{video_name}.mp4', fourcc, float(FPS), frameSize)

for i in range(number_of_img*(number_of_step-1)):
  im = pilimg.open(f'/content/Cartoon-StyleGan2/asset/{outdir}/out-{i}.png')
  np_img = np.asarray(im)

  frame = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
  video.write(frame)
video.release()

#--------------------------------------
# Show a Interpolation Video in Colab
#--------------------------------------

from IPython.display import HTML
from base64 import b64encode
mp4 = open(f'/content/Cartoon-StyleGan2/asset/{outdir}/{video_name}.mp4','rb').read()

data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""

""" % data_url)

 

2. Style CLIP

: 텍스트 명령어를 이용해 이미지를 원하는 방식으로 수정할 수 있는 모델

 

[Example 1]

# 이미지 프로세스 매개변수

• seed : 잠재 벡터를 생성하는데 사용되는 랜덤 seed 값
• description : 이미지 텍스트 설명 
• ckpt1,2 : ffhq, naver webtoon 데이터셋으로 사전 훈련된 stylegan2 모델의 가중치가 저장된 ff 체크포인트 파일
• style_size : 픽셀 단위의 모델 해상도
• latent_dim : 잠재 벡터 차원
• lr_rampup : 학습 속도 증가 조절
• lr : 학습률(초기 학습 속도)
• step : 최적화 단계 수
• l2_lambda : l2 regularization에 사용되는 정규화 매개변수 -> 특징벡터 정규화 과정에서 사용됨
• latent_path : 잠재 벡터 경로(랜덤 시드를 사용하지 않는 경우)
• truncation : 특정 latent 벡터가 평균 latent 벡터로부터 벗어나지 벡터를 잘라내줌(약 0.7정도에서 좋은 이미지 생성)
• save_intermediate_image_every : 최적화 프로세스 동안 중간 이미지를 저장하는 빈도
• device : 최적화에 사용되는 장치
• results_dir : 최종 이미지 저장 디렉토리

description = 'a really angry face' #@param {type:"string"}

seed = 536357 #@param {type:"slider", min:0, max:1000000, step:1}

latent_path = None #@param {type:"string"}

optimization_steps = 300 #@param {type:"number"}

l2_lambda = 0.004 #@param {type:"number"}

create_video = False #@param {type:"boolean"}

result_dir = "asset/results1" #@param {type:"string"}

device = 'cuda'

# -----------------------------
args = {
    "seed" : seed,
    "description": description,
    "ckpt": "/content/Cartoon-StyleGan2/networks/ffhq256.pt",
    "ckpt2": "/content/Cartoon-StyleGan2/networks/NaverWebtoon.pt",
    "stylegan_size": 256,
    "latent_dim" : 14,
    "lr_rampup": 0.05,
    "lr": 0.1,
    "step": optimization_steps,
    "l2_lambda": l2_lambda,
    "latent_path": latent_path,
    "truncation": 0.7,
    "save_intermediate_image_every": 1 if create_video else 20,
    "device" : "cuda",
    "results_dir": result_dir,
}

 

• args에 정의된 입력값으로 run_optimization.py 실행
• run_optimization.py는 텍스트 설명과 일치하는 이미지 생성 모델의 최적화 수행의 main 함수
• 최종 생성 이미지, 초기 및 최종 잠재 코드를 반환

from run_optimization import main
from argparse import Namespace

final_result, latent_init2, latent_fin2 = main(Namespace(**args))

 

• ffhq, metface 데이터셋으로 훈련된 generator1, generator2 로드
• 생성할 이미지 크기 256x256, 잠재 벡터 차원 512, 레이어 8
• 생성기 로드 후 무작위의 4096개 샘플을 평균 잠재 벡터로 만들어 각 generator의 절단 값 잠재벡터 초기화 
• latent1_init1, lantent_fin1 간의 보간을 수행한 후 생성된 이미지 표시
  • 보간 단계 수 5, 보간의 강도 1.5(값이 클수록 잠재 벡터 사이에 뚜렷한 변화 나타남).
  • swap = True : 두 개의 생성기에서의 스타일 혼합기술 적용, 2번쨰 레이어에서 스타일 혼합 수행

import torch

from model import Generator
from utils import imshow, tensor2image

device = 'cuda'

# ---------------
# Generator
# ---------------

# Generator1
network1='/content/Cartoon-StyleGan2/networks/ffhq256.pt' 
network1 = torch.load(network1)

generator1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
generator1.load_state_dict(network1["g_ema"], strict=False)

trunc1 = generator1.mean_latent(4096)

# Generator2
Target_network = "Metface_StructureLoss" #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network2=f'/content/Cartoon-StyleGan2/networks/{Target_network}.pt' 
network2 = torch.load(network2)

generator2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
generator2.load_state_dict(network2["g_ema"], strict=False)

trunc2 = generator2.mean_latent(4096)

# ---------------
# Interpolation
# ---------------

number_of_step = 5 #@param {type:"slider", min:0, max:30, step:1}
latent_interp = torch.zeros(number_of_step, latent_init1.shape[1], latent_init1.shape[2]).to(device)
strength = 1.5 #@param {type:"slider", min:1, max:3, step:0.25}
swap = True #@param {type:"boolean"}
swap_layer_num = 2 #@param {type:"slider", min:1, max:6, step:1}


with torch.no_grad():
    for j in range(number_of_step):

        latent_interp[j] = latent_init1 + strength * (latent_fin1-latent_init1) * float(j/(number_of_step-1))

        imgs_gen1, save_swap_layer = generator1([latent_interp],
                                input_is_latent=True,                                     
                                truncation=0.7,
                                truncation_latent=trunc1,
                                swap=swap, swap_layer_num=swap_layer_num,
                                randomize_noise=False)
        imgs_gen2, _ = generator2([latent_interp],
                                input_is_latent=True,                                     
                                truncation=0.7,
                                swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
                                truncation_latent=trunc2)

imshow(tensor2image(torch.cat([img_gen for img_gen in imgs_gen1], dim=2)), 20)
imshow(tensor2image(torch.cat([img_gen for img_gen in imgs_gen2], dim=2)), 20)

imgs_gen1

imgs_gen2

의문점 1) trunc1,2가 4096개의 샘플로부터 만들어낸 임의의 평균 잠재 벡터가 맞나?

truncation 변수는 잘림 정도값을 설정해주어 이미지 품질과 다양성의 균형을 맞춰주게 되는데  trunc1,2는 무슨 역할을 하는 것인가?

의문점 2) latent_init1, latent_final1 이 어떤 이미지에 대한 잠재벡터인가?(run_optimization.py 코드를 봐야할 것 같다.)

의문점 3) 텍스트 명령어로 이미지를 변경할 수 있는 과정이 아직 잘 이해가 안됨 

 

 

3. Style Mixing 

: 두 이미지의 스타일 벡터를 혼합하여 이미지 생성

 

# generator 로드, 잠재 벡터 보간

• ffhq, roamce 101 generator 생성
• .pt 파일의 사전 훈련된 모델 가중치로 generator1,2 초기화 
• latent1, latent2 잠재 벡터 초기화
• swap, save_swap_layer변수로 스타일 혼합 수행, 혼합 정보 저장 레이어 지정
• 두 개의 generator로 부터 생성된 두 이미지 사이를 보간

import torch

from model import Generator
from utils import imshow, tensor2image

device = 'cuda'
Target_network = "Romance101" #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
swap = False #@param {type:"boolean"}
save_swap_layer = 1 #@param {type:"slider", min:1, max:6, step:1}

# ---------------
# Generator
# ---------------

# Generator1
network1='/content/Cartoon-StyleGan2/networks/ffhq256.pt' 
network1 = torch.load(network1)

generator1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
generator1.load_state_dict(network1["g_ema"], strict=False)

trunc1 = generator1.mean_latent(4096)

# latent1
seed1 = 627356 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed1)
latent1 = torch.randn(1, 14, 512, device=device)
latent1 = generator1.get_latent(latent1)


# Generator2
network2=f'/content/Cartoon-StyleGan2/networks/{Target_network}.pt' 
network2 = torch.load(network2)

generator2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
generator2.load_state_dict(network2["g_ema"], strict=False)

trunc2 = generator2.mean_latent(4096)

# latent2
seed2 = 159972 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed2)
latent2 = torch.randn(1, 14, 512, device=device)
latent2 = generator2.get_latent(latent2)

# ---------------
# Interpolation
# ---------------

number_of_step = 6 #@param {type:"slider", min:0, max:10, step:1}
latent_interp = torch.zeros(number_of_step, latent1.shape[1], latent1.shape[2]).to(device)

with torch.no_grad():
    for j in range(number_of_step):

        latent_interp[j] = latent1 + (latent2-latent1) * float(j/(number_of_step-1))

        imgs_gen1, save_extract_layer = generator1([latent_interp],
                                input_is_latent=True,                                     
                                truncation=0.7,
                                truncation_latent=trunc1,
                                swap=swap, swap_layer_num=swap_layer_num,
                                )
                                
        imgs_gen2, _ = generator2([latent_interp],
                                input_is_latent=True,                                     
                                truncation=0.7,
                                truncation_latent=trunc2,
                                swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
                                )

imshow(tensor2image(torch.cat([img_gen for img_gen in imgs_gen1], dim=2)), 20)
imshow(tensor2image(torch.cat([img_gen for img_gen in imgs_gen2], dim=2)), 20)

💡 두 개의 이미지를 보간하는 코드가 이해가 잘 안됐는데 chat GPT가 꽤나 자세히 설명해줌...!

#InterPolation 코드는 ffhq 이미지에 대한 latent1, Romance101이미지에 대한 latent2 벡터를 선형으로 보간하고 보간한 벡터를 generator 1, 2에 입력하여 중간 이미지를 생성한다. 

number_of_step = 6은 두 개의 입력 이미지 사이에 생성할 중간 이미지의 수 6개이다. 

보간된 잠재 벡터를 저장할 latent_interp 를 0으로 초기화된 (6, 14, 512) 크기의 텐서로 초기화 한다. 

latent_interp[j] = latent1 + (latent2-latent1) * float(j/(number_of_step-1)) 공식은 latent1, lantent2를 선형 보간하여 잠재벡터를 생성하는 공식이다. 

6단계를 거치면서 두 벡터 사이의 선형 보간을 계산하여 잠재 보간 벡터가 생성된다. 

그런 다음 중간 이미지를 생성하기 위해 generator1,2에 보간된 잠재 벡터를 사용하여 이미지를 만들어낸다.

generator1 은 network1(ffhq) 입력 이미지에 해당하는 중간 이미지를 생성하고 generator2는 network2(romance101) 이미지에 해당하는 중간 이미지를 생성하는데 사용된다.

generator1과 generator2는 모두 디스크에서 로드된 각각의 사전 훈련된 모델을 사용하여 코드의 초기에 초기화되었고

latent_interp 텐서는 두 generator에 대한 입력으로 사용되며, input_is_latent =True로 설정되어 생성기에 대한 입력이 이미지 텐서가 아닌 잠재 벡터임을 나타낸다. 

imshow 함수는 생성된 중간 이미지를 각 생성기마다 하나씩 두 개의 개별 창에 표시하는 데 사용된다. 

(어렵다 ㅎ.ㅎ)

 

 

• latent1_mixing, lantent2_mixing(혼합 비율)을 사용해서 latent3 벡터 생성
• 3개의 벡터를 연결해서 최종 잠재 벡터 생성하여 ffhq, romance101 이미지 생성

from utils import imshow, tensor2image

device = 'cuda'
truncation = 0.7 

target_network ="Romance101" #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
swap = False #@param {type:"boolean"}
swap_layer_num = 1 #@param {type:"slider", min:1, max:6, step:1}
# ----------------------
# Source Images (FFHQ)
# ----------------------

network1='/content/Cartoon-StyleGan2/networks/ffhq256.pt' 
network1 = torch.load(network1)

# Genearator1
g1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g1.load_state_dict(network1["g"], strict=False)
trunc1 = g1.mean_latent(4096)

# latent1
seed1 = 124322 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed1)
r_latent1 = torch.randn(1, 14, 512, device=device)
latent1 = g1.get_latent(r_latent1)

# latent2
seed2 = 398781 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed2)
r_latent2 = torch.randn(1, 14, 512, device=device)
latent2 = g1.get_latent(r_latent2)

# latent mixing
latent_mixing1 = 4 #@param {type:"slider", min:0, max:15, step:1}
latent_mixing2 = 9 #@param {type:"slider", min:0, max:15, step:1}

latent3 = torch.cat([latent1[:,:latent_mixing1,:], latent2[:,latent_mixing1:latent_mixing2,:], latent1[:,latent_mixing2:,:]], dim = 1)
# latent4 = torch.cat([latent1[:,:,:latent_mixing2], latent2[:,:,latent_mixing2:]], dim = 2)

# Latent !
latent = torch.cat([latent1, latent2, latent3], dim = 0)

# generate image
img1, save_swap_layer = g1(
    [latent],
    truncation=truncation,
    truncation_latent=trunc1,
    swap=swap, swap_layer_num=swap_layer_num,
    input_is_latent=True,
)

# =================================================

# ----------------------
# Target Images (Cartoon)
# ----------------------


network2 =f'/content/Cartoon-StyleGan2/networks/{target_network}.pt'
network2 = torch.load(network2)

# Genearator2
g2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g2.load_state_dict(network2["g"], strict=False)
trunc2 = g2.mean_latent(4096)

# generate image

img2, _ = g2(
    [latent],
    input_is_latent=True,
    truncation=truncation,
    truncation_latent=trunc2,
    swap=swap, swap_layer_num=swap_layer_num, swap_layer_tensor=save_swap_layer,
)

# =================================================

# Show Image

ffhq = torch.cat([img1[0], img1[1], img1[2]], dim=2)
cartoon = torch.cat([img2[0], img2[1], img2[2]], dim=2)

imshow(tensor2image(torch.cat([ffhq, cartoon], dim = 1)), 15)

3. make a video

• ffhq, romance101 이미지 생성 

#@title  { run: "auto" }

# Save Image
import torch
import os
from utils import make_image
from torchvision.utils import save_image, make_grid
import PIL.Image as pilimg
from model import Generator

device = 'cuda'
truncation = 0.7 

target_network = 'Romance101' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']

# ----------------------
# Source Images (FFHQ)
# ----------------------


network1='/content/Cartoon-StyleGan2/networks/ffhq256.pt' 
network1 = torch.load(network1)

# Genearator1
g1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g1.load_state_dict(network1["g"], strict=False)
trunc1 = g1.mean_latent(4096)

# latent1
seed1 = 627356 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed1)
r_latent1 = torch.randn(1, 14, 512, device=device)
latent1 = g1.get_latent(r_latent1)

# latent2
seed2 = 159972 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed2)
r_latent2 = torch.randn(1, 14, 512, device=device)
latent2 = g1.get_latent(r_latent2)

# latent mixing
latent_mixing1 = 1 #@param {type:"slider", min:0, max:14, step:1}
latent_mixing2 = 419 #@param {type:"slider", min:0, max:512, step:1}
latent3 = torch.cat([latent1[:,:latent_mixing1,:], latent2[:,latent_mixing1:,:]], dim = 1)
latent4 = torch.cat([latent1[:,:,:latent_mixing2], latent2[:,:,latent_mixing2:]], dim = 2)

# Latent !
latent = torch.cat([latent1, latent2, latent3, latent4], dim = 0)


# ----------------------
# Target Images (Cartoon)
# ----------------------

network2=f'/content/Cartoon-StyleGan2/networks/{target_network}.pt'
network2 = torch.load(network2)

# Genearator2
g2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g2.load_state_dict(network2["g"], strict=False)
trunc2 = g2.mean_latent(4096)

 

• 두 소스 이미지(ffhq, romance101)를 ffhq 이미지를 배경으로 삼아 스타일을 혼합하여 이미지 생성
• 3개의 잠재벡터를 연결하여 최종 잠재 벡터를 만들고 혼합이미지 생성
• 생성된 이미지 png 파일로 저장

#@markdown Latent Mixing1

if not os.path.isdir('./asset/latent_mix1'):
   os.makedirs('./asset/latent_mix1', exist_ok=True)

for latent_mixing1 in range(15):
    latent3 = torch.cat([latent1[:,:latent_mixing1,:], latent2[:,latent_mixing1:,:]], dim = 1)
    
    latent = torch.cat([latent2, latent1, latent3], dim = 0)

    img1, mix_layer = g1(
        [latent],
        input_is_latent=True,
        truncation=truncation,
        truncation_latent=trunc1,
        extract_layer=True, extract_layer_num=1,
    )
    img2, _ = g2(
        [latent],
        input_is_latent=True,
        truncation=truncation,
        truncation_latent=trunc2,
        extract_layer=False, extract_layer_num=1, put_layer = mix_layer,
    )

    img_name = "/content/Cartoon-StyleGan2/asset/latent_mix1/latent_mixing1_" + str(latent_mixing1) + ".png"
    grid = make_grid(torch.cat([img1, img2], dim=0),
                    nrow=3,
                    normalize=True,
                    range=(-1,1))

    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = pilimg.fromarray(ndarr)
    im.save(img_name)

 

• 잠재벡터를 15개의 벡터로 분할하여 각각의 벡터에 대한 이미지 생성 후 비디오로 결합 

#@title  { run: "auto" }


# =====================================

from PIL import Image
import cv2 
from torchvision.utils import save_image
import PIL.Image as pilimg
import numpy as np

img_name = "/content/Cartoon-StyleGan2/asset/latent_mix1/latent_mixing1_0.png"
image = Image.open(img_name)

frameSize = image.size
FPS = 3 #@param {type:"slider", min:0, max:10, step:1}
fourcc = cv2.VideoWriter_fourcc(*'XVID')

video = cv2.VideoWriter('/content/Cartoon-StyleGan2/asset/latent_mix1/latent_mixing1.mp4', fourcc, float(FPS), frameSize)

for i in range(15):

    im = pilimg.open("/content/Cartoon-StyleGan2/asset/latent_mix1/latent_mixing1_" + str(i) + ".png")
    np_img = np.asarray(im)

    frame = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
    video.write(frame)

video.release()
#--------------------------
# Show a Interpolation Video in Colab
#--------------------------

from IPython.display import HTML
from base64 import b64encode
mp4 = open('/content/Cartoon-StyleGan2/asset/latent_mix1/latent_mixing1.mp4','rb').read()

data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""

""" % data_url)

 

• generator1, generator2로 부터 생성된 이미지의 잠재 벡터 간의 혼합을 보여주는 이미지 그리드 생성

from utils import imshow, tensor2image

device = 'cuda'
truncation = 0.7 

save_latent = False #@param {type:"boolean"}
target_network ="NaverWebtoon" #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
n_sample = 6 #@param {type:"slider", min:1, max:10, step:1}
# latent mixing
latent_mixing1 = 6 #@param {type:"slider", min:0, max:14, step:1}


# ----------------------
# Source Images (FFHQ)
# ----------------------

network1=f'/content/Cartoon-StyleGan2/networks/{target_network}.pt' 
network1 = torch.load(network1)

# Genearator1
g1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g1.load_state_dict(network1["g"], strict=False)
trunc1 = g1.mean_latent(4096)

# latent1
if not save_latent:
  r_latent1 = torch.randn(n_sample, 14, 512, device=device)
  latent1 = g1.get_latent(r_latent1)

# generate image
img1, _ = g1(
    [latent1],
    truncation=truncation,
    truncation_latent=trunc1,
    input_is_latent=True,
)

images = [torch.cat([img for img in img1], 2)]


for i in range(n_sample):
  latent = latent1[i].unsqueeze(0).repeat(n_sample, 1, 1)
  latent2 = torch.cat([latent[:,:latent_mixing1,:], latent1[:,latent_mixing1:,:]], dim = 1)
  
  # generate image
  img1, _ = g1(
      [latent2],
      truncation=truncation,
      truncation_latent=trunc1,
      input_is_latent=True,
  )

  images.append(torch.cat([img for img in img1], 2))

imshow(tensor2image(torch.cat([img for img in images], 1)), 20)

 

 

4. closed-form factorization

• 로컬 시스템에서 작성된 코드라 내 환경에서는 실행할 수 없는 코드라고 뜬다..!
• apply_factor.py 파일 로드
• stylegan2-ada를 실행하여 생성된 이미지 집합에 대한 사전 훈련된 모델에 factor를 적용한 후 비디오로 출력하는 코드

!python apply_factor.py --index=7 --degree=14 --seed=123177 --n_sample=5 \
                        --ckpt='/content/Cartoon-StyleGan2/networks/ffhq256.pt' --ckpt2='/content/Cartoon-StyleGan2/networks/Romance101.pt' \
                        --factor='/content/Cartoon-StyleGan2/networks/factor.pt' --outdir='asset/results-sefa' --video 

from IPython.display import Image
Image(open('/content/Cartoon-StyleGan2/asset/results-sefa/video.gif','rb').read())

 

• SEFA + styleGAN2-ADA를 적용하여 사전 훈련된 두 모델 사이에 보간된 이미지 생성
• 잠재 방향의 인덱스와 보간 정도를 지정할 수 잇음
• 출력 이미지 디렉토리 저장, 비디오 형식 가능

#@title { display-mode: "form", run: "auto"}

import torch
import cv2 
import numpy as np
from torchvision.utils import save_image, make_grid
import PIL.Image as pilimg

from model import Generator


device='cuda'
n_sample=5
truncation = 0.7
trunc = 0.7

# Eigen-Vector
factor='/content/Cartoon-StyleGan2/factor.pt'
eigvec = torch.load(factor)["eigvec"].to(device)


# =============================================

# Genearaotr1
network1 = 'ffhq256' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network1=f'/content/Cartoon-StyleGan2/networks/{network1}.pt' 
network1 = torch.load(network1)

g1 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g1.load_state_dict(network1["g_ema"], strict=False)


# Generator2
network2 = 'TrueBeauty' #@param ['ffhq256', 'NaverWebtoon', 'NaverWebtoon_StructureLoss', 'NaverWebtoon_FreezeSG', 'Romance101', 'TrueBeauty', 'Disney', 'Disney_StructureLoss', 'Disney_FreezeSG', 'Metface_StructureLoss', 'Metface_FreezeSG']
network2=f'/content/Cartoon-StyleGan2/networks/{network2}.pt' 

network2 = torch.load(network2)

g2 = Generator(256, 512, 8, channel_multiplier=2).to(device)
g2.load_state_dict(network2["g_ema"], strict=False)

# latent
seed = 177838 #@param {type:"slider", min:0, max:1000000, step:1}
torch.manual_seed(seed)
latent = torch.randn(n_sample, 14, 512, device=device)
latent = g1.get_latent(latent)

# latent direction & scalar
index=19 #@param {type:"slider", min:0, max:512, step:1}
degree=15 #@param {type:"slider", min:0, max:30, step:1}


# =============================================

# directory to save image
outdir = 'result-sefa' #@param {type:"string"}
if not os.path.isdir(f'{outdir}'):
   os.makedirs(f'./asset/{outdir}', exist_ok=True)

imgs = []

for deg in range(degree):

  direction = 0.5 * deg * eigvec[:, index].unsqueeze(0) #torch.Size([1, 512])

  img1, _ = g1(
      [latent + direction],
      truncation=truncation,
      truncation_latent=trunc,
      input_is_latent=True,
  )

  img2, _ = g2(
      [latent + direction],
      truncation=truncation,
      truncation_latent=trunc,
      input_is_latent=True,
  )


  grid = make_grid(torch.cat([img1, img2], 0),
                   nrow=n_sample,
                   normalize=True,
                   range=(-1,1),
                   )
  ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
  im = pilimg.fromarray(ndarr)
  im.save(f'/content/Cartoon-StyleGan2/asset/{outdir}/out-{deg}.png')