%tensorflow_version 1.x
!nvidia-smi

from google.colab import drive
from pathlib import Path

content_path = Path('/').absolute() / 'content'
drive_path = content_path / 'drive'
drive.mount(str(drive_path))

stylegan2_repo_url  = 'https://github.com/ArthurFDLR/stylegan2-ada' # Original repository: https://github.com/NVlabs/stylegan2-ada
project_path        = drive_path / 'MyDrive' / 'StyleGAN2-ADA'
stylegan2_repo_path = project_path / 'stylegan2-ada'

# Create project folder if inexistant
if not project_path.is_dir():
    %mkdir "{project_path}"
%cd "{project_path}"

for dir in ['in', 'out', 'datasets', 'training']:
    if not (project_path / dir).is_dir():
        %mkdir {dir}
if not (project_path / 'datasets' / 'source').is_dir():
    %mkdir "{project_path / 'datasets' / 'source'}"

# Download StyleGAN2-ada
!git config --global user.name "ArthurFDLR"
!git config --global user.email "[email protected]"
if stylegan2_repo_path.is_dir():
    !git -C "{stylegan2_repo_path}" fetch origin
    !git -C "{stylegan2_repo_path}" checkout origin/main -- *.py
else:
    print("Install StyleGAN2-ADA")
    !git clone {stylegan2_repo_url}

dataset_name = 'NightSky'
datasets_source_path = project_path / 'datasets' / 'source' / (dataset_name + '.zip')
if datasets_source_path.is_dir():
    print("Dataset ready for import.")
else:
    print('Upload your images dataset as {}'.format(datasets_source_path))

local_dataset_path = content_path / 'dataset'
if not local_dataset_path.is_dir():
    print("Importing dataset...")
    %mkdir "{local_dataset_path}"
    %cp -a "{project_path / 'datasets' / 'source' / (dataset_name + '.zip')}" "{local_dataset_path}"
    print("Zip file succesfuly imported")
else:
    print('Zip file allready imported')

import zipfile
with zipfile.ZipFile(str(local_dataset_path / (dataset_name + '.zip')), 'r') as zip_ref:
    zip_ref.extractall(str(local_dataset_path))
print('Extraction completed')

local_images_path = local_dataset_path / 'images'
local_dataset_path /= 'tfr'

if (local_dataset_path).is_dir():
    print('\N{Heavy Exclamation Mark Symbol} Dataset already created \N{Heavy Exclamation Mark Symbol}')
    print('Delete current dataset folder ({}) to regenerate tfrecords.'.format(local_dataset_path))
else:
    %mkdir "{local_dataset_path}"
    !python "{stylegan2_repo_path / 'dataset_tool.py'}" create_from_images \
        "{local_dataset_path}" "{local_images_path}"

!python "{stylegan2_repo_path / 'train.py'}" --help

training_path = project_path / 'training' / dataset_name
if not training_path.is_dir():
    %mkdir "{training_path}"

#how often should the model generate samples and a .pkl file
snapshot_count = 2
#should the images be mirrored left to right?
mirrored = True
#should the images be mirrored top to bottom?
mirroredY = False
#metrics? 
metric_list = None
#augments
augs = 'bgc'

resume_from = 'ffhq1024'

!python "{stylegan2_repo_path / 'train.py'}" --outdir="{training_path}" \
    --data="{local_dataset_path}" --resume="{resume_from}" \
    --snap={snapshot_count} --augpipe={augs} \
    --mirror={mirrored} --mirrory={mirroredY} \
    --metrics={metric_list} #--dry-run

%pip install opensimplex
!python "{stylegan2_repo_path / 'generate.py'}" generate-images --help 

from numpy import random
seed_init = random.randint(10000)
nbr_images = 6

generation_from = 'ffhq1024'

!python "{stylegan2_repo_path / 'generate.py'}" generate-images \
    --outdir="{project_path / 'out'}" --trunc=0.7 \
    --seeds={seed_init}-{seed_init+nbr_images-1} --create-grid \
    --network={generation_from}

%pip install opensimplex
!python "{stylegan2_repo_path / 'generate.py'}" generate-latent-walk --help 

from numpy import random
walk_types = ['line', 'sphere', 'noiseloop', 'circularloop']
latent_walk_path = project_path / 'out' / 'latent_walk'
if not latent_walk_path.is_dir():
    %mkdir "{latent_walk_path}"

explored_network = 'ffhq1024'

seeds = random.randint(10000) for  i in range(50)
print("Base seeds:", seeds)
!python "{stylegan2_repo_path / 'generate.py'}" generate-latent-walk --network="{explored_network}" \
    --outdir="{latent_walk_path}" --trunc=0.7 --walk-type="{walk_types[2]}" \
    --seeds={','.join(map(str, seeds))} --frames {len(seeds)*20}