from time import time
from os.path import join
import os
import sys
import lpips
from GAN_hessian_compute import hessian_compute, get_full_hessian
from torchvision.transforms import ToPILImage
from torchvision.utils import make_grid
from GAN_utils import loadBigGAN, loadStyleGAN2, BigGAN_wrapper, StyleGAN2_wrapper, loadStyleGAN, StyleGAN_wrapper
from hessian_analysis_tools import plot_spectra, compute_hess_corr
from lpips import LPIPS
ImDist = LPIPS(net="squeeze")
datadir = r"E:\OneDrive - Washington University in St. Louis\HessNetArchit\StyleGAN"
os.makedirs(datadir, exist_ok=True)
#%%
StyleGAN = loadStyleGAN()
SD = StyleGAN.state_dict()
#%%
shuffled_SD = {}
for name, Weight in SD.items():
idx = torch.randperm(Weight.numel())
W_shuf = Weight.view(-1)[idx].view(Weight.shape)
shuffled_SD[name] = W_shuf
#%%
torch.save(shuffled_SD, join(datadir, "StyleGAN_shuffle.pt"))
# print(name, Weight.shape, Weight.mean().item(), Weight.std().item())
#%%
StyleGAN_sf = loadStyleGAN()
StyleGAN_sf.load_state_dict(torch.load(join(datadir, "StyleGAN_shuffle.pt")))
G_sf = StyleGAN_wrapper(StyleGAN_sf)
#%%
#%%
vis_eigen_frame(evc_avg,
eva_avg,
PG,
figdir=figdir,
namestr="PGG_orig",
page_B=15,
eig_rng=(0, 120),
maxdist=22,
rown=7,
sphere=False,
transpose=False,
RND=0)
#%% StyleGAN
SGGAN = loadStyleGAN() # , size=256
SG = StyleGAN_wrapper(SGGAN)
figdir = join(axesdir, "StyleGAN")
os.makedirs(figdir, exist_ok=True)
data = np.load(join(rootdir, "StyleGAN", "H_avg_%s.npz" % "StyleGAN"))
H_avg, eva_avg, evc_avg = data["H_avg"], data["eva_avg"], data["evc_avg"]
#%%
RND = np.random.randint(1E4)
ref_codes = np.random.randn(8, 512)
for eigidx in range(40):
vis_eigen_action(evc_avg[:, -eigidx - 1],
ref_codes,
SG,
figdir=figdir,
namestr="SG_eig%d_lin" % (eigidx + 1),
maxdist=2.5,
vis_eigen_explore(np.zeros((1,512)), evc_avg, eva_avg, SG2, figdir=figsumdir, namestr="StyleGAN2_%s_orig_expon"%modelsnm, transpose=False,
eiglist=list(range(0,20)), maxdist=.75E-3, rown=5, sphere=False, scaling=eva_avg[-1:-21:-1]**(-1/3))
#%% Linear exploration from origin
vis_eigen_explore(np.zeros((1,512)), evc_avg, eva_avg, SG2, figdir=figsumdir, namestr="StyleGAN2_%s_orig_expon"%modelsnm, transpose=False,
eiglist=list(range(0,20)), maxdist=.4E-3, rown=5, sphere=False, scaling=eva_avg[-1:-21:-1]**(-1/3))
#%%
refvec = np.random.randn(1,512)
vis_eigen_explore(refvec, evc_avg, eva_avg, SG2, figdir=figsumdir, namestr="StyleGAN2_%s_sphexpon"%modelsnm, transpose=False, eiglist=list(range(0,20)), maxdist=0.25, rown=5, sphere=True, scaling=eva_avg[-1:-21:-1]**(-1/3))
#%% StyleGAN
"""StyleGAN for Faces"""
Hessdir = join(rootdir, 'StyleGAN')
modelnm = "StyleGAN"
SGAN = loadStyleGAN()
SG = StyleGAN_wrapper(SGAN)
with np.load(join(Hessdir, "H_avg_%s.npz"%modelnm)) as data:
eva_avg, evc_avg, = data["eva_avg"], data["evc_avg"], # feat_col = data["feat_col"]
#%%
refvec = torch.randn(1, 512).cuda()
img = SG.visualize(refvec)
ToPILImage()(make_grid(img).cpu()).show()
#%%
# vis_eigen_frame(evc_avg, eva_avg, SG, ref_code=refvec, figdir=figsumdir, namestr="StyleGAN", transpose=False,
# eiglist=None, eig_rng=(0, 20), maxdist=1, rown=5, sphere=False, )
#%% Linear exploration from origin
_=vis_eigen_explore(np.zeros((1, 512)), evc_avg, eva_avg, SG, figdir=figsumdir, namestr="StyleGAN_orig_expon", transpose=False,
eiglist=list(range(0,20)), maxdist=.8E-3, rown=5, sphere=False, scaling=eva_avg[-1:-21:-1]**(-1/2.5))
#%%
_=vis_eigen_explore(np.zeros((1, 512)), evc_avg, eva_avg, SG, figdir=figsumdir, namestr="StyleGAN_orig_expon", transpose=False,
12,
], :],
G=BBG,
maxdist=2,
rown=5,
transpose=False,
namestr="BigBiGAN_norm9_eig%d" % eigid,
figdir=figdir)
#%% StyleGAN2
from GAN_hessian_compute import hessian_compute
from GAN_utils import loadStyleGAN, StyleGAN_wrapper
figdir = r"E:\OneDrive - Washington University in St. Louis\Hessian_summary\StyleGAN2"
#%% Cats
modelname = "stylegan2-cat-config-f"
npzdir = r"E:\Cluster_Backup\StyleGAN2\stylegan2-cat-config-f"
SGAN = loadStyleGAN(modelname + ".pt", size=256, channel_multiplier=2) #
G = StyleGAN_wrapper(SGAN)
eigval_col, eigvec_col, feat_col, meta = scan_hess_npz(
npzdir,
npzpat="Hess_BP_(\d*).npz",
evakey='eva_BP',
evckey='evc_BP',
featkey="feat")
feat_arr = np.array(feat_col).squeeze()
#%%
eigid = 5
mtg = vis_eigen_action(eigvec=eigvec_col[0][:, -eigid - 1],
ref_codes=feat_arr[[
0,
2,
4,