%reload_ext autoreload
%autoreload 2

from fastai import *
from fastai.vision import *

#path = Path('data/lsun')
#path_img = path/'bedroom'
path = untar_data(URLs.PETS)

train_ds = (ImageItemList.from_folder(path).label_const(0., label_cls=FloatList)
               .transform(tfms=[crop_pad(size=64, row_pct=(0,1), col_pct=(0,1))],size=64))

data = (ImageDataBunch.create(train_ds, valid_ds=None, path=path)
                     .normalize(stats = [torch.tensor([0.5,0.5,0.5]), torch.tensor([0.5,0.5,0.5])]))

wgan = models.BasicGAN(64, 100, 3, 64, 1)
learn = Learner(data, wgan, loss_func=NoopLoss(), opt_func=optim.RMSprop, wd=0.)

cb = callbacks.GANTrainer(learn, bs=128, n_disc_iter=callbacks.first_disc_iter)
learn.callbacks.append(cb)
learn.fit(1, 1e-4)

learn.recorder.losses[:10]

cb.dlosses[:10]

%debug

x,y = next(iter(learn.data.train_dl))

tst = learn.model(callbacks.create_noise(x,64,100,False), gen=True)
img = data.denorm(tst[0].cpu()).numpy().clip(0,1).transpose(1,2,0)
plt.imshow(img)

learn.model(tst)

learn.model(x)