FaceGenerator.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "cmd_line_run = False\n",
    "if not cmd_line_run:\n",
    "    %matplotlib inline\n",
    "collab_mode = False\n",
    "\n",
    "if collab_mode and not cmd_line_run:\n",
    "    # set up tensorflow in collab\n",
    "    %tensorflow_version 2.x\n",
    "# imports\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import warnings # This ignore all the warning messages\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from os import path\n",
    "import os\n",
    "import time\n",
    "\n",
    "print(\"Tensorflow version is\", tf.__version__, \", device name\", tf.test.gpu_device_name())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def allow_memory_growth():\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            # Currently, memory growth needs to be the same across GPUs\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "            logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n",
    "            print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
    "        except RuntimeError as e:\n",
    "            # Memory growth must be set before GPUs have been initialized\n",
    "            print(e)\n",
    "\n",
    "# allow_memory_growth()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def get_time():\n",
    "    return time.strftime(\"%d-%m-%Y_%H-%M-%S\")\n",
    "\n",
    "def read_epoch_index(settings):\n",
    "    if os.path.exists(settings.epoch_filename):\n",
    "        with open(settings.epoch_filename, 'r') as f:\n",
    "            try:\n",
    "                return int(f.readline())\n",
    "            except:\n",
    "                return 0\n",
    "    return 0\n",
    "\n",
    "def save_epoch_index(settings, epoch_index):\n",
    "    with open(settings.epoch_filename, 'w') as f:\n",
    "        f.write(str(epoch_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "metadata": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def process_image(img, image_shape):\n",
    "    img = tf.cast(img, tf.float32)/127.5-1 # IMPORTANT, image's pixels are in the range <-1, 1>\n",
    "    # crop 128x128, then resize it\n",
    "    to_keep = 150\n",
    "    y_offset, x_offset = (img.shape[0] - to_keep) // 2, (img.shape[1] - to_keep) // 2\n",
    "    img = img[y_offset:y_offset+to_keep, x_offset:x_offset+to_keep, :]\n",
    "    img = tf.image.resize(img, image_shape)\n",
    "    return img\n",
    "\n",
    "def load_image(filename):\n",
    "    img = tf.io.read_file(filename)\n",
    "    img = tf.image.decode_jpeg(img)\n",
    "    return img\n",
    "\n",
    "def convert_from_output_to_image(images):\n",
    "    # converts <-1, 1> to <0,1>\n",
    "    return tf.clip_by_value((images+1)/2, 0, 1)\n",
    "\n",
    "def display_image_from_dataset(data):\n",
    "    if cmd_line_run:\n",
    "        return\n",
    "    for batch in data.take(1):\n",
    "        image, attributes = batch\n",
    "        img_ = convert_from_output_to_image(image[0])\n",
    "        plt.imshow(img_)\n",
    "        print(img_.shape, np.min(img_), np.max(img_))\n",
    "\n",
    "def save_generated_image(settings, epoch, size):\n",
    "    save_dir = settings.generated_images_path\n",
    "    if not path.exists(save_dir):\n",
    "        os.makedirs(save_dir)\n",
    "    name = path.join(save_dir,\n",
    "                     'img_{}_{}_{}.png'.format(size, epoch, get_time()))\n",
    "    plt.savefig(name)\n",
    "\n",
    "\n",
    "def show_images(images, epoch, settings, save_images=False, display_images=False):\n",
    "    print(\"image pixels range\", np.min(images), np.max(images), \"std\", np.std(images))\n",
    "    num_of_images = min(10, images.shape[0])\n",
    "    # (x, y=1)\n",
    "    plt.figure(figsize=(num_of_images, 1))\n",
    "    for i in range(num_of_images):\n",
    "        plt.subplot(1, num_of_images, i + 1)\n",
    "        img = images[i, :, :, :].numpy() #\n",
    "        # img = convert_from_output_to_image(img) # images are already converted in gen_step\n",
    "        # img = (img * 127.5 + 127.5).astype(np.uint8)\n",
    "        plt.imshow(img)\n",
    "        plt.axis('off')\n",
    "    \n",
    "\n",
    "    if save_images:\n",
    "        save_generated_image(settings, epoch, size=images.shape[1])\n",
    "    if display_images and not cmd_line_run:\n",
    "        plt.show()\n",
    "\n",
    "def load_dataset(dataset_path, image_shape, use_manually_downloaded_dataset=True, preprocess_images=True, shuffle_size=500, seed=101):\n",
    "    if use_manually_downloaded_dataset:\n",
    "        img_path = path.join(dataset_path, '*.jpg')\n",
    "        data = tf.data.Dataset.list_files(img_path, seed=seed)\\\n",
    "                              .shuffle(shuffle_size)\\\n",
    "                              .map(load_image)\n",
    "        # for each image return a tuple (image, attributes), in this case we have no attributes\n",
    "        if preprocess_images:\n",
    "            data = data.map(lambda x: (process_image(x, image_shape), dict()))\n",
    "        else:\n",
    "            data = data.map(lambda x: (x, dict()))\n",
    "            \n",
    "    else:\n",
    "        dataset_name = 'celeb_a'\n",
    "        data = tfds.load(dataset_name, split=tfds.Split.TRAIN)\\\n",
    "                   .shuffle(shuffle_size)\n",
    "        # for each image return a tuple (image, attributes), ignore 'landmarks'\n",
    "        if preprocess_images:\n",
    "            data = data\\\n",
    "                .map(lambda x: (process_image(x['image'], image_shape), x['attributes']))\n",
    "        else:\n",
    "            data = data\\\n",
    "                .map(lambda x: (x['image'], x['attributes']))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class Settings:\n",
    "    def __init__(self, collab_mode):\n",
    "        self.root_local_path = os.getcwd()\n",
    "        self.root_gdrive_path = '/content/drive'\n",
    "        self.gdrive_project_path = 'My Drive/pp/GSN/FaceGenerator'\n",
    "        self.dataset_name = \"celeb_a\"\n",
    "        self.subdataset_dir=\"1000\"\n",
    "        self.dataset_image_size = (28, 28)\n",
    "        self.image_size = (128, 128)\n",
    "        assert self.image_size in [(64, 64), (128, 128), (256, 256)], 'image_size must be on of 64x64, 128x128, 256x256'\n",
    "        self.image_channels = 3\n",
    "        self.generator_input_shape = (self.image_size[0], self.image_channels)\n",
    "        self.gdrive_mounted = False\n",
    "        self.collab_mode = collab_mode\n",
    "        self.batch_size = 20\n",
    "        self.epochs = 100\n",
    "        self.save_models = False # save models at the end?\n",
    "        self.mount_gdrive()\n",
    "        # variables from pytorch\n",
    "        self.init_w = 4 # ?\n",
    "        self.init_h = 4 #?\n",
    "        self.gen_init_filters = 64 # how many filters*2 (channels) should be in the first layer of the generator, n_g number in the paper\n",
    "        self.df_dim = 64\n",
    "        self.gf_dim=self.gen_init_filters # ?\n",
    "        self.z_dim=(100, 100) # size of the vector with normal distribution noise\n",
    "        self.embeddings_shape = (100, 128)\n",
    "        self.n_g = self.gen_init_filters\n",
    "        self.n_d = 64\n",
    "        self.is_conditional=True\n",
    "        self.unconditional_loss_coefficient = 1.0\n",
    "        self.branch_num = 1 # 1 == 64; 2 == 64, 128; 3 == 64, 128, 256\n",
    "        assert 1 <= self.branch_num <= 3, 'branch num in range <1, 3>'\n",
    "        \n",
    "    @property\n",
    "    def run_name(self):\n",
    "        return 'stack_gan_embedding_from_attr_small_{}'.format(self.branch_num)\n",
    "        # return \"{}_epochs_{}_batch_{}\".format(self.epochs, self.batch_size, self.subdataset_dir)\n",
    "    \n",
    "    @property\n",
    "    def image_shape(self):\n",
    "        return (*self.image_size, self.image_channels)\n",
    "\n",
    "    @property\n",
    "    def download_path(self):\n",
    "        return path.join('.', 'datasets', self.dataset_name)\n",
    "    # def download_path(self):\n",
    "    #     return path.join(self.get_base_path, 'datasets', self.dataset_name)\n",
    "\n",
    "    @property\n",
    "    def dataset_path(self):\n",
    "        return path.join(self.download_path, self.subdataset_dir)\n",
    "\n",
    "    @property\n",
    "    def tensorboard_log_dir(self):\n",
    "        return path.join(self.get_base_path, 'saved_state', self.run_name, 'tensorboard_logs')\n",
    "\n",
    "    @property\n",
    "    def checkpoint_dir(self):\n",
    "        return path.join(self.get_base_path, 'saved_state', self.run_name, \"ckpt\")\n",
    "\n",
    "    @property\n",
    "    def model_save_path(self):\n",
    "        return path.join(self.get_base_path, 'saved_state', self.run_name, 'models')\n",
    "    \n",
    "    @property\n",
    "    def generated_images_path(self):\n",
    "        return path.join(self.get_base_path, 'saved_state', self.run_name, 'generated_images')\n",
    "    \n",
    "    @property\n",
    "    def epoch_filename(self):\n",
    "        return path.join(self.get_base_path, 'saved_state', self.run_name, 'epoch_number')\n",
    "    \n",
    "    @property\n",
    "    def get_base_path(self):\n",
    "        if self.collab_mode:\n",
    "            return path.join(self.root_gdrive_path, self.gdrive_project_path)\n",
    "        else:\n",
    "            return self.root_local_path\n",
    "\n",
    "    def mount_gdrive(self):\n",
    "        if self.collab_mode:\n",
    "            if cmd_line_run:\n",
    "                print('cmd line run, abort mounting gdrive')\n",
    "                exit(-1)\n",
    "            from google.colab import drive\n",
    "            project_path = path.join(self.root_gdrive_path, self.gdrive_project_path)\n",
    "            self.gdrive_project_path = path.join(self.root_gdrive_path, self.gdrive_project_path)\n",
    "            drive.mount(self.root_gdrive_path)\n",
    "            self.gdrive_mounted = True\n",
    "        \n",
    "            path_with_imports = path.join(self.root_gdrive_path, self.gdrive_project_path)\n",
    "            print(\"Files in path\", path_with_imports)\n",
    "            !ls /content/drive/My\\ Drive/pp/GSN/FaceGenerator\n",
    "            if path_with_imports not in os.sys.path:\n",
    "                os.sys.path.append(path_with_imports)\n",
    "                \n",
    "class DatasetCache:\n",
    "    def __init__(self, use_manually_downloaded_dataset):\n",
    "        self.path = \"\"\n",
    "        self.batch_size = 0\n",
    "        self._data = None\n",
    "        # should it use dataset downloaded manually or the one downloaded by tfds\n",
    "        self.use_manually_downloaded_dataset = use_manually_downloaded_dataset\n",
    "\n",
    "    def load_data(self, settings):\n",
    "        if self._data is not None and self.path == settings.dataset_path:\n",
    "            self.batch_size = settings.batch_size\n",
    "            return self.data\n",
    "        else:\n",
    "            print(\"downloading and loading data\")\n",
    "            if self.use_manually_downloaded_dataset:\n",
    "                self.download_dataset(settings)\n",
    "            self._data = load_dataset(\n",
    "                settings.dataset_path,\n",
    "                settings.image_size,\n",
    "                use_manually_downloaded_dataset=self.use_manually_downloaded_dataset\n",
    "            )\n",
    "            self.data_path = settings.dataset_path\n",
    "            self.batch_size = settings.batch_size\n",
    "            return self.data\n",
    "        \n",
    "    @property\n",
    "    def data(self):\n",
    "        if self._data is None:\n",
    "            return None\n",
    "        else:\n",
    "            return self._data.batch(self.batch_size)\n",
    "\n",
    "    def download_dataset(self, settings):\n",
    "        import dataset_helpers as ds_helpers\n",
    "        '''Downloads data to dataset_path/dataset_name directory'''\n",
    "        print('dataset download path is {}'.format(settings.download_path))\n",
    "        ds_helpers.download_extract('celeba', settings.download_path)\n",
    "        \n",
    "class TensorboardManager():\n",
    "    def __init__(self):\n",
    "        self.log_path = ''\n",
    "        self.train_summary_writer = None\n",
    "        self.test_summary_writer = None\n",
    "        \n",
    "    def initialize(self, settings):\n",
    "        should_be_updated = False\n",
    "        if settings.collab_mode and self.log_path != \"tensorboard_logs\":\n",
    "            should_be_updated = True\n",
    "            self.log_path = \"/content/drive/My Drive/pp/GSN/FaceGenerator/saved_state/run_img_align_celeba/tensorboard_logs\"\n",
    "        elif not settings.collab_mode and self.log_path != settings.tensorboard_log_dir:\n",
    "            should_be_updated = True\n",
    "            self.log_path = settings.tensorboard_log_dir\n",
    "                \n",
    "        if should_be_updated:\n",
    "            self.train_summary_writer = tf.summary.create_file_writer(path.join(self.log_path, 'train'))\n",
    "            self.test_summary_writer = tf.summary.create_file_writer(path.join(self.log_path, 'test'))\n",
    "            print('Initialized tensorboard log dir with path', self.log_path)\n",
    "            self.launch(settings.collab_mode)\n",
    "        \n",
    "    def launch(self, collab_mode):\n",
    "        if collab_mode:\n",
    "            if cmd_line_run:\n",
    "                print('cmd line mode, abort launching tensorboard')\n",
    "                returns\n",
    "            %reload_ext tensorboard\n",
    "            %tensorboard --logdir \"/content/drive/My Drive/pp/GSN/FaceGenerator/saved_state/run_img_align_celeba/tensorboard_logs\"\n",
    "            from tensorboard import notebook\n",
    "            notebook.list() # View open TensorBoard instances\n",
    "        else:\n",
    "            print('open tensorboard with command')\n",
    "            print('tensorboard --bind_all --logdir {}'.format(self.log_path))\n",
    "            \n",
    "    def save(self, run_type, data, description, timestamp, datatype, step):\n",
    "        '''\n",
    "        run_type: either `train` or `test`\n",
    "        data: represents scalar, images or list of scalars\n",
    "        description: should be of a length of data, i.e. if data is a scalar, description should be a string\n",
    "        datatype: one of 'scalar', 'scalars', 'images'\n",
    "        '''\n",
    "        def _save(writer):\n",
    "            with writer.as_default():\n",
    "                if datatype == 'scalars':\n",
    "                    for value, name in zip(data, description):\n",
    "                        tf.summary.scalar('{}_{}'.format(name, timestamp), value, step=step)\n",
    "                elif datatype == 'scalar':\n",
    "                    tf.summary.scalar('{}_{}'.format(description, timestamp), data, step=step)\n",
    "                elif datatype == 'images':\n",
    "                    tf.summary.image('{}_{}'.format(description, timestamp), data, step=step)\n",
    "                else:\n",
    "                    print('unknown type', datatype)\n",
    "\n",
    "        if run_type == 'train' and self.train_summary_writer:\n",
    "            _save(self.train_summary_writer)\n",
    "        elif run_type == 'test' and self.train_summary_writer:\n",
    "            _save(self.test_summary_writer)\n",
    "        else:\n",
    "            print('unrecognized option `run_type` or selected writer', run_type,'is None')    \n",
    "            \n",
    "class Environment():\n",
    "    def __init__(self, collab_mode):\n",
    "        self.settings = Settings(collab_mode)\n",
    "        self.models = dict()\n",
    "        self.datasetCache = DatasetCache(use_manually_downloaded_dataset=False)\n",
    "        self.checkpointManager = None\n",
    "        self.tensorboard = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_time():\n",
    "    return time.strftime(\"%d-%m-%Y_%H-%M-%S\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "# test downloading and loading data\n",
    "\n",
    "dataset = DatasetCache(use_manually_downloaded_dataset=False)\n",
    "settings = Settings(collab_mode)\n",
    "if collab_mode:\n",
    "    settings.subdataset_dir = 'img_align_celeba'\n",
    "dataset.load_data(settings)\n",
    "\n",
    "if not collab_mode:\n",
    "    display_image_from_dataset(dataset.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "# load network for embeddings creation\n",
    "embedding_net_name = 'small_embedding_10attr_encoder' # encoder_out outputs attributes\n",
    "with open('./embedding/{}_structure.json'.format(embedding_net_name), 'r') as f:\n",
    "    embedding_net = tf.keras.models.model_from_json(f.read())\n",
    "embedding_net.summary()\n",
    "embedding_net.load_weights('./embedding/{}_weights.h5'.format(embedding_net_name))\n",
    "print('loaded network for embeddings:', embedding_net_name)\n",
    "\n",
    "# use this one if using embedding network that takes some subset of attributes as an input\n",
    "# the subset must be compatible with the one that the embedding network was trained on\n",
    "attributes_subset = [\n",
    "        'Attractive', 'Bald', 'Bangs',\n",
    "        'Big_Nose', 'Heavy_Makeup', 'Blond_Hair',\n",
    "        'Male', 'Smiling', 'Young', 'Straight_Hair'\n",
    "]\n",
    "# use this one if using embedding network that takes an image as an input\n",
    "# attributes_subset = None\n",
    "\n",
    "# def prepare_embeddings_using_net(batch):\n",
    "#     imgs, attributes = batch\n",
    "#     if attributes_subset is not None:\n",
    "#         attributes = {key: attributes[key] for key in attributes if key in attributes_subset}\n",
    "#     embeddings = tf.stack(list(attributes.values()))\n",
    "    \n",
    "#     imgs = tf.image.resize(imgs, (64, 64))\n",
    "#     return imgs, embedding_net.predict(imgs) # TODO <0,1>\n",
    "\n",
    "\n",
    "def prepare_embeddings_from_batch(batch, output_type=tf.int8, attributes_subset = None, use_embedding_net=True):\n",
    "    imgs, attributes = batch\n",
    "    if not use_embedding_net:\n",
    "        return imgs, embeddings\n",
    "    \n",
    "    if attributes_subset is not None:\n",
    "        attributes = {key: attributes[key] for key in attributes if key in attributes_subset}\n",
    "        embeddings = tf.stack(list(attributes.values()))\n",
    "        embeddings = tf.transpose(embeddings)\n",
    "        embeddings = tf.cast(embeddings, dtype=output_type)\n",
    "        return imgs, embedding_net.predict(embeddings)\n",
    "    else:\n",
    "        imgs, _ = batch\n",
    "        return imgs, embedding_net.predict(tf.image.resize(imgs, (64, 64)))\n",
    "\n",
    "    \n",
    "def describe_embedding(orig_attributes, example):\n",
    "    image, embedding_vector = example\n",
    "    # display image\n",
    "    img_ = (image+1)/2\n",
    "    plt.imshow(img_)\n",
    "    # describe image\n",
    "    for key, value in zip(orig_attributes.keys(), embedding_vector):\n",
    "        print(key, ' '*(30-len(key)), value.numpy())\n",
    "\n",
    "def test_creating_embedding(dataset):\n",
    "    for batch in dataset.data.take(1): # take one batch\n",
    "        # get imgs and embeddings from that batch\n",
    "        imgs, embeddings = prepare_embeddings_from_batch(batch,\n",
    "                                                         output_type=tf.int8, \n",
    "                                                         attributes_subset=attributes_subset,\n",
    "                                                         use_embedding_net=True)\n",
    "        print(embeddings.shape)\n",
    "        # describe first image in the batch and display it\n",
    "        # describe_embedding(batch[1], (imgs[0], embeddings[0]) )\n",
    "    \n",
    "# imgs, embeddings = prepare_embeddings_from_batch(batch, settings)\n",
    "test_creating_embedding(dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "display_image_from_dataset(dataset.data)\n",
    "if not collab_mode and not cmd_line_run:\n",
    "    !ls \"./datasets/celeb_a/img_align_celeba\" | wc -l"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Noise Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S6YpknpB6a0C",
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class NoiseGenerator(tf.keras.layers.Layer):\n",
    "\n",
    "    def __init__(self, distribution_size):\n",
    "        super().__init__()\n",
    "        self.distribution_size = distribution_size\n",
    "        # self.data_distributions = self.add_weight(shape=(num_classes, distribution_size), trainable=True)\n",
    "        # self.data_distributions = tf.tile(tf.range(0, num_classes, dtype=tf.float32)[:, tf.newaxis], [1, distribution_size])\n",
    "        # TODO:\n",
    "\n",
    "    def call(self, batch_size):\n",
    "        # dists = tf.nn.embedding_lookup(self.data_distributions, inputs)\n",
    "        # dists += tf.random.uniform(tf.shape(dists), -0.35, 0.35)\n",
    "        # return dists\n",
    "        # TODO\n",
    "        return tf.random.normal([batch_size, self.distribution_size, 3])\n",
    "    \n",
    "    def diverse_distributions_loss(self):\n",
    "        # TODO\n",
    "        return None\n",
    "    \n",
    "    \n",
    "class StackGANNoiseGenerator(tf.keras.layers.Layer):\n",
    "    def __init__(self, distribution_size=100):\n",
    "        super().__init__()\n",
    "        self.distribution_size = distribution_size\n",
    "\n",
    "    def call(self, inputs):\n",
    "        # batch_size x distribution_size\n",
    "        return tf.random.normal([tf.shape(inputs)[0], self.distribution_size])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1xwRLPZYmJtL",
    "pycharm": {}
   },
   "source": [
    "Loss functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TnfRHsNBmLsd",
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "bce = tf.keras.losses.BinaryCrossentropy()\n",
    "kl = tf.keras.losses.KLDivergence()\n",
    "\n",
    "def min_max_discriminator_loss(real_out, gen_out):\n",
    "    real_loss = bce(tf.ones_like(real_out), real_out)\n",
    "    gen_loss = bce(tf.zeros_like(gen_out), gen_out)\n",
    "    return real_loss + gen_loss\n",
    "\n",
    "def min_max_generator_loss(gen_out):\n",
    "    return - min_max_discriminator_loss(tf.ones_like(gen_out), gen_out)\n",
    "\n",
    "def stack_gan_discriminator_loss(real_out, fake_out):\n",
    "    return min_max_discriminator_loss(real_out, fake_out)\n",
    "\n",
    "def stack_gan_generator_loss(gen_out):\n",
    "    return bce(tf.ones_like(gen_out), gen_out)\n",
    "\n",
    "def w_discriminator_loss(real_out, gen_out):\n",
    "    res = - (tf.reduce_mean(real_out) - tf.reduce_mean(gen_out))\n",
    "    return res\n",
    "\n",
    "def w_generator_loss(gen_out):\n",
    "    return - tf.reduce_mean(gen_out)\n",
    "\n",
    "def kl_generator_loss(real, gen):\n",
    "    return kl(real, gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {}
   },
   "source": [
    "Layers for generator and discriminator from StackGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class GLU(tf.keras.layers.Layer):\n",
    "    '''\n",
    "    Multiplies first half of the last dimension\n",
    "    by value of the second half of the last dimension.\n",
    "    Layer introduced in https://github.com/hanzhanggit/StackGAN-v2/blob/master/code/model.py#L47\n",
    "    \n",
    "    Returns the same tensor with its last dimension halved and multiplied by sigmoid function activation.\n",
    "    '''\n",
    "    def __init__(self, **kwargs):\n",
    "        super(GLU, self).__init__(**kwargs)\n",
    "\n",
    "    def call(self, data):\n",
    "        rank = data.shape.rank\n",
    "        tensor_shape = data.shape.as_list()\n",
    "        last_dim_half = tensor_shape[-1]//2\n",
    "        # automatic slicing is impossible since the first dimension is None (batch size)\n",
    "        \n",
    "\n",
    "        # begins_1, begins_2 = np.zeros((rank,), dtype=np.int32), np.zeros((rank,), dtype=np.int32)\n",
    "        # begins_2[-1] = last_dim_half # set offset for the last dimension\n",
    "\n",
    "        # sizes = [tensor_shape[i] for i in range(rank)]\n",
    "        # sizes[-1] = last_dim_half # get half of elements from the last dimension\n",
    "\n",
    "        # return tf.slice(data, begins_1, sizes) * tf.math.sigmoid(tf.slice(data, begins_2, sizes))\n",
    "        \n",
    "        \n",
    "        if rank == 2:\n",
    "            return data[:, :last_dim_half] * tf.math.sigmoid(data[:, last_dim_half:])\n",
    "        elif rank == 3:\n",
    "            return data[:, :, :last_dim_half] * tf.math.sigmoid(data[:, :, last_dim_half:])\n",
    "        elif rank == 4:\n",
    "            return data[:, :, :, :last_dim_half] * tf.math.sigmoid(data[:, :, :, last_dim_half:])\n",
    "        elif rank == 5:\n",
    "            return data[:, :, :, :, :last_dim_half] * tf.math.sigmoid(data[:, :, :, :, last_dim_half:])\n",
    "\n",
    "class GeneratorResidualLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, filters, index, **kwargs):\n",
    "        super(GeneratorResidualLayer, self).__init__(**kwargs)\n",
    "        \n",
    "        self.block = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(filters*2, 3, 1, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            GLU(),\n",
    "            tf.keras.layers.Conv2D(filters, 3, 1, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "        ], name='generator_residual_layer_{}'.format(index))\n",
    "\n",
    "    def call(self, data):\n",
    "        residual = data\n",
    "        out = self.block(data)\n",
    "        return tf.keras.layers.add([out, residual])\n",
    "    \n",
    "class GeneratorUpsampleLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, filters, **kwargs):\n",
    "        super(GeneratorUpsampleLayer, self).__init__(**kwargs)\n",
    "        self.block = tf.keras.Sequential([\n",
    "            tf.keras.layers.UpSampling2D(),\n",
    "            tf.keras.layers.Conv2D(filters*2, 3, 1, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            GLU()\n",
    "        ], name='generator_upsample_layer')\n",
    "\n",
    "    def call(self, data):\n",
    "        return self.block(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FwrjmA4t6Oxn",
    "pycharm": {}
   },
   "source": [
    "Generator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class GANNetwork(tf.keras.Model):\n",
    "    def __init__(self, model_name=\"Network\", **kwargs):\n",
    "        self.is_conditional = False\n",
    "        if 'is_conditional' in kwargs:\n",
    "            self.is_conditional = kwargs['is_conditional']\n",
    "            # remove that entry from kwargs\n",
    "            # tf raises an exception when a keyword is not from allowed_kwargs set\n",
    "            del kwargs['is_conditional']\n",
    "        super().__init__(name=model_name, **kwargs)\n",
    "        \n",
    "        self.model = None\n",
    "        \n",
    "\n",
    "    def print_layers(self):\n",
    "        print(self.model)\n",
    "        for layer in self.model:\n",
    "            print(layer.name, \":\", layer.input_shape, \"->\", layer.output_shape)\n",
    "\n",
    "    def summary(self):\n",
    "        self.model.summary()\n",
    "        \n",
    "    @property\n",
    "    def output_layer(self):\n",
    "        return GANNetwork.get_last_layer(self.model)\n",
    "    \n",
    "    @property\n",
    "    def output_shape(self):\n",
    "        return GANNetwork.get_last_layer_output_shape(self.model)\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, data, training):\n",
    "        return self.model(data)\n",
    "\n",
    "    def save_model(self, save_path):\n",
    "        if not path.exists(save_path):\n",
    "            os.makedirs(save_path)\n",
    "        filename = path.join(save_path, '{}_{}'.format(self.name, get_time()))\n",
    "        print(\"Saving model\", self.name, \"as\", filename)\n",
    "        self.model.save(filename)\n",
    "        \n",
    "    def compile_model(self, optimizer, loss):\n",
    "        self.model.compile(optimizer=optimizer, loss=loss)\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_last_layer(model):\n",
    "        return model.get_layer(index=-1)\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_last_layer_output_shape(model):\n",
    "        layer = GANNetwork.get_last_layer(model)\n",
    "        if layer is not None:\n",
    "            return layer.output_shape[1:] # don't take batch_size into consideration\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {},
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class GeneratorInitStage(GANNetwork):\n",
    "    #input_shape = cfg.GAN.Z_DIM + cfg.GAN.EMBEDDING_DIM\n",
    "    # or \n",
    "    # input_shape = cfg.GAN.Z_DIM if without conditioning\n",
    "    def __init__(self, noise_shape, embeddings_shape, filters, settings, model_name=\"GeneratorInitStage\", **kwargs):\n",
    "        # filters should be `gen_init_filters`\n",
    "        super().__init__(model_name, **kwargs)\n",
    "        self.initial_factor = 16\n",
    "        self.filters = self.initial_factor*filters\n",
    "        # if is conditional then the second dimension is has size of concatenated noise and embedding\n",
    "        \n",
    "        self.in_shape = \\\n",
    "            (noise_shape[1] + embeddings_shape[1]) if self.is_conditional else noise_shape[1]\n",
    "#         print('in shape', self.in_shape)\n",
    "        \n",
    "        self.model = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(settings.init_w*settings.init_h*self.filters*2, input_shape=(self.in_shape, )),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            GLU(),\n",
    "            tf.keras.layers.Reshape([settings.init_w, settings.init_h, self.filters]),\n",
    "            GeneratorUpsampleLayer(self.filters//2),\n",
    "            GeneratorUpsampleLayer(self.filters//4),\n",
    "            GeneratorUpsampleLayer(self.filters//8),\n",
    "            GeneratorUpsampleLayer(self.filters//16)\n",
    "        ], name='generator_intermediate_init_output')\n",
    "        \n",
    "        self.output_model = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(3, 3, 1, 'same', input_shape=self.output_shape)\n",
    "        ], name = 'generator_init_stage_output')\n",
    "        \n",
    "    @property\n",
    "    def image_shape(self):\n",
    "        return GANNetwork.get_last_layer_output_shape(self.output_model)\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self, data, training):\n",
    "        intermediate_output = self.model(data, training=training)\n",
    "        output = self.output_model(intermediate_output, training=training)\n",
    "        return intermediate_output, output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "def prepare_embedding_tensor(embedding, output_shape):\n",
    "    '''\n",
    "    Extends embedding tensor to match output shape.\n",
    "    [batch, z] -> [batch, x, y, z], where x, y == output_shape\n",
    "    '''\n",
    "    embedding_shape = embedding.shape.as_list()\n",
    "    repeat_factor = tf.math.reduce_prod(output_shape)\n",
    "    res = tf.keras.backend.repeat(embedding, repeat_factor)\n",
    "    reshape_s = [embedding_shape[0], *output_shape, embedding_shape[-1]]\n",
    "    return tf.reshape(res, reshape_s)\n",
    "\n",
    "\n",
    "class GeneratorNextStage(GANNetwork):    \n",
    "    @staticmethod\n",
    "    def concat_embedding_and_previous_output(embedding, previous_output):\n",
    "        repeat_factor = 1\n",
    "#         if previous_output.shape.rank == 4:\n",
    "#             # get second and third dimension\n",
    "#             previous_output_shape = previous_output.shape.as_list()[1:3]\n",
    "#         else:\n",
    "#             raise(\"Rank of the previous_output is not equal to 4\")\n",
    "        embedding = prepare_embedding_tensor(embedding, (64, 64))\n",
    "        return tf.concat([embedding, previous_output], axis=-1)\n",
    "    \n",
    "\n",
    "    def __init__(self, embeddings_shape, number_of_residuals, previous_stage_shape, filters, model_name=\"GeneratorNextStage\", **kwargs):\n",
    "        super().__init__(model_name, **kwargs)\n",
    "        \n",
    "        # size should be equal to\n",
    "        # (init_stage_output_y, init_stage_output_x, init_stage_output_channels\n",
    "        self.in_shape = previous_stage_shape\n",
    "        if self.is_conditional:\n",
    "            # concat embeddings with channel\n",
    "            self.in_shape = (self.in_shape[0], self.in_shape[1], self.in_shape[2] + embeddings_shape[1])\n",
    "        \n",
    "        self.model = tf.keras.Sequential()\n",
    "        \n",
    "        self.model.add(tf.keras.layers.Dense(filters*2, input_shape=self.in_shape))\n",
    "        self.model.add(tf.keras.layers.BatchNormalization())\n",
    "        self.model.add(GLU())\n",
    "        \n",
    "        self.model.add(GeneratorResidualLayer(filters, index=0))\n",
    "        for i in range(number_of_residuals-1):\n",
    "            self.model.add(GeneratorResidualLayer(filters, i+1))\n",
    "        \n",
    "        self.model.add(GeneratorUpsampleLayer(filters//2))\n",
    "        \n",
    "        self.output_model = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(3, 3, 1, 'same', input_shape=self.output_shape)\n",
    "        ], name = 'generator_next_stage_output')\n",
    "        \n",
    "    @property\n",
    "    def image_shape(self):\n",
    "        return GANNetwork.get_last_layer_output_shape(self.output_model)\n",
    "        \n",
    "    @tf.function\n",
    "    def call(self, data, training):\n",
    "        intermediate_output = self.model(data, training=training)\n",
    "        output = self.output_model(intermediate_output, training=training)\n",
    "        return intermediate_output, output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8N1JmZB06SKO",
    "pycharm": {}
   },
   "source": [
    "Discriminator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class DiscriminatorEncodeImageLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, filters, **kwargs):\n",
    "        super(DiscriminatorEncodeImageLayer, self).__init__(**kwargs)\n",
    "        self.block = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(filters, 4, 2, 'same', use_bias=False),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2),\n",
    "            \n",
    "            tf.keras.layers.Conv2D(filters*2, 4, 2, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2),\n",
    "            \n",
    "            tf.keras.layers.Conv2D(filters*4, 4, 2, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2),\n",
    "            \n",
    "            \n",
    "            tf.keras.layers.Conv2D(filters*8, 4, 2, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2),\n",
    "        ], name='discriminator_encoding_layer')\n",
    "\n",
    "    def call(self, data):\n",
    "        return self.block(data)\n",
    "    \n",
    "class DiscriminatorDownBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, filters, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.block = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(filters, 4, 2, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2)\n",
    "        ])\n",
    "        \n",
    "    def call(self, data):\n",
    "        return self.block(data)\n",
    "    \n",
    "class Dicriminator3x3LeakyReLU(tf.keras.layers.Layer):\n",
    "    def __init__(self, filters, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.block = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(filters, 3, 1, 'same', use_bias=False),\n",
    "            tf.keras.layers.BatchNormalization(),\n",
    "            tf.keras.layers.LeakyReLU(alpha=.2)\n",
    "        ])\n",
    "        \n",
    "    def call(self, data):\n",
    "        return self.block(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class StackGANDiscriminator(GANNetwork):\n",
    "    def __init__(self, input_shape, filters, embeddings_shape, output_image_size, model_name=\"StackGANDiscriminator\", **kwargs):\n",
    "        super().__init__(model_name, **kwargs)\n",
    "        \n",
    "        self.filters = filters\n",
    "        \n",
    "        # define intermediate models\n",
    "        self.model = tf.keras.Sequential([\n",
    "            DiscriminatorEncodeImageLayer(filters = self.filters, input_shape=input_shape)\n",
    "        ])\n",
    "        if output_image_size == 128:\n",
    "            self.model.add(DiscriminatorDownBlock(filters=self.filters*16))\n",
    "            self.model.add(Dicriminator3x3LeakyReLU(filters=self.filters*8))\n",
    "\n",
    "        elif output_image_size == 256:\n",
    "            self.model.add(DiscriminatorDownBlock(filters=self.filters*16))\n",
    "            self.model.add(DiscriminatorDownBlock(filters=self.filters*32))\n",
    "            self.model.add(Dicriminator3x3LeakyReLU(filters=self.filters*16))\n",
    "            self.model.add(Dicriminator3x3LeakyReLU(filters=self.filters*8))\n",
    "\n",
    "        self.conditional_model = tf.keras.Sequential()\n",
    "        self.unconditional_model = tf.keras.Sequential([\n",
    "            tf.keras.layers.Conv2D(1, kernel_size=4, strides=4, input_shape=self.output_shape)\n",
    "        ])\n",
    "        \n",
    "        if self.is_conditional:\n",
    "            conditional_input = (self.output_shape[0], self.output_shape[1], self.output_shape[2]+embeddings_shape[1])\n",
    "            self.conditional_model.add(Dicriminator3x3LeakyReLU(filters=self.filters*8, input_shape=conditional_input))\n",
    "            self.conditional_model.add(tf.keras.layers.Conv2D(1, kernel_size=4, strides=4))    \n",
    "            \n",
    "    @tf.function\n",
    "    def call(self, data, training):\n",
    "        if self.is_conditional:\n",
    "            image, embedding = data # if unconditional then embedding == None\n",
    "        else:\n",
    "            image = data\n",
    "            \n",
    "        encoding_output = self.model(image, training=training)\n",
    "\n",
    "        unconditional_output = self.unconditional_model(encoding_output, training=training)\n",
    "        unconditional_loss = tf.math.sigmoid(tf.reshape(unconditional_output, [-1, 1]))\n",
    "        \n",
    "        if self.is_conditional:\n",
    "            # join embedding and image\n",
    "            encode_layer_shape = self.model.output_shape[1:3] # get y and x of the model output shape\n",
    "            embedding = prepare_embedding_tensor(embedding, encode_layer_shape)\n",
    "            conditional_input = tf.keras.layers.concatenate([embedding, encoding_output], axis=-1)\n",
    "            \n",
    "            conditional_output = self.conditional_model(conditional_input, training=training)\n",
    "            conditional_loss = tf.math.sigmoid(tf.reshape(conditional_output, [-1, 1]))\n",
    "            \n",
    "            return conditional_loss, unconditional_loss\n",
    "        \n",
    "        return None, unconditional_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "enYw9PQ_6h_o",
    "pycharm": {}
   },
   "source": [
    "Training step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class ModelWrapper:\n",
    "    def __init__(self, model_object, model_optimizer_fun, model_loss_fun):\n",
    "        self.model = model_object\n",
    "        self.optimizer_fun = model_optimizer_fun\n",
    "        self.loss_fun = model_loss_fun\n",
    "        self.loss = 0.0\n",
    "\n",
    "\n",
    "def stack_gan_train_step_template(generators, discriminators, noise_generator, settings):\n",
    "    '''\n",
    "    generators: {'name': ModelWrapper}\n",
    "    discriminators: {'name': ModelWrapper}\n",
    "    noise_generator: noise generator model\n",
    "    '''\n",
    "    training = True\n",
    "    unconditional_loss_coefficient = settings.unconditional_loss_coefficient\n",
    "    is_conditional = settings.is_conditional\n",
    "    branch_num = settings.branch_num\n",
    "    \n",
    "    @tf.function\n",
    "    def _train_step_template(images, embeddings): # to be executed before\n",
    "        # assuming images are 128x128\n",
    "        imgs_64 = tf.image.resize(images, (64, 64))\n",
    "        \n",
    "        if branch_num > 2:\n",
    "            imgs_256 = tf.image.resize(images, (256, 256))\n",
    "        \n",
    "        gen_trainable_variables = generators['64'].model.trainable_variables\n",
    "        discr_trainable_variables = discriminators['64'].model.trainable_variables\n",
    "\n",
    "        losses_g = []\n",
    "        losses_d = []\n",
    "        with tf.GradientTape() as d_tape, tf.GradientTape() as g_tape:\n",
    "                \n",
    "            noise = noise_generator(tf.shape(images)[0])\n",
    "\n",
    "            ### 64 ###\n",
    "            gen_64_input = tf.concat([embeddings, noise], axis=1)\n",
    "            gen_64_out, gen_64_image_out = generators['64'].model(gen_64_input, training=training)\n",
    "            \n",
    "            # discriminator\n",
    "            fake_discr_64_output = discriminators['64'].model([gen_64_image_out, embeddings], training=training)\n",
    "            real_discr_64_output = discriminators['64'].model([imgs_64, embeddings], training=training)\n",
    "            \n",
    "            # calculate generator loss (conditional + unconditional*unconditional_coeff)\n",
    "            discriminators['64'].loss = discriminators['64']\\\n",
    "                .loss_fun(real_discr_64_output[0], fake_discr_64_output[0])\n",
    "            \n",
    "            discriminators['64'].loss += discriminators['64']\\\n",
    "                .loss_fun(real_discr_64_output[1], fake_discr_64_output[1])*unconditional_loss_coefficient\n",
    "            \n",
    "            # calculate generator loss (conditional + unconditional*unconditional_coeff)\n",
    "            generators['64'].loss = generators['64'].loss_fun(fake_discr_64_output[0])\n",
    "            generators['64'].loss += generators['64'].loss_fun(fake_discr_64_output[1])*unconditional_loss_coefficient\n",
    "            \n",
    "            losses_g.append(generators['64'].loss)\n",
    "            losses_d.append(discriminators['64'].loss)\n",
    "            \n",
    "            if branch_num > 1:\n",
    "                ### 128 ###\n",
    "                gen_128_input = GeneratorNextStage.concat_embedding_and_previous_output(\n",
    "                        embedding=embeddings,\n",
    "                        previous_output=gen_64_out)\n",
    "                gen_128_out, gen_128_image_out = generators['128'].model(gen_128_input, training=training)\n",
    "\n",
    "                # discriminator\n",
    "                fake_discr_128_output = discriminators['128'].model([gen_128_image_out, embeddings], training=training)\n",
    "                real_discr_128_output = discriminators['128'].model([images, embeddings], training=training)\n",
    "\n",
    "                discriminators['128'].loss = discriminators['128']\\\n",
    "                    .loss_fun(real_discr_128_output[0], fake_discr_128_output[0])\n",
    "                discriminators['128'].loss += discriminators['128']\\\n",
    "                    .loss_fun(real_discr_128_output[1], fake_discr_128_output[1])*unconditional_loss_coefficient\n",
    "\n",
    "                # generator loss\n",
    "                generators['128'].loss = generators['128'].loss_fun(fake_discr_128_output[0])\n",
    "                generators['128'].loss += generators['128'].loss_fun(fake_discr_128_output[1])*unconditional_loss_coefficient\n",
    "\n",
    "                gen_trainable_variables = generators['64'].model.trainable_variables + generators['128'].model.trainable_variables\n",
    "                discr_trainable_variables += discriminators['128'].model.trainable_variables\n",
    "                losses_g.append(generators['128'].loss)\n",
    "                losses_d.append(discriminators['128'].loss)\n",
    "\n",
    "            if branch_num > 2:\n",
    "                ### 256 ###\n",
    "                gen_256_input = GeneratorNextStage.concat_embedding_and_previous_output(\n",
    "                        embedding=embeddings,\n",
    "                        previous_output=gen_128_out)\n",
    "                gen_256_out, gen_256_image_out = generators['256'].model(gen_256_input, training=training)\n",
    "\n",
    "                # discriminator\n",
    "                fake_discr_256_output = discriminators['256'].model([gen_256_image_out, embeddings], training=training)\n",
    "                real_discr_256_output = discriminators['256'].model([imgs_256, embeddings], training=training)\n",
    "                \n",
    "                discriminators['256'].loss = discriminators['256']\\\n",
    "                    .loss_fun(real_discr_256_output[0], fake_discr_256_output[0])\n",
    "                discriminators['256'].loss += discriminators['256']\\\n",
    "                    .loss_fun(real_discr_256_output[1], fake_discr_256_output[1])*unconditional_loss_coefficient\n",
    "                # generator loss\n",
    "                generators['256'].loss = generators['256'].loss_fun(fake_discr_256_output[0])\n",
    "                generators['256'].loss += generators['256'].loss_fun(fake_discr_256_output[1])*unconditional_loss_coefficient\n",
    "                    \n",
    "                gen_trainable_variables += generators['256'].model.trainable_variables\n",
    "                discr_trainable_variables += discriminators['256'].model.trainable_variables\n",
    "                losses_g.append(generators['256'].loss)\n",
    "                losses_d.append(discriminators['256'].loss)\n",
    "\n",
    "        \n",
    "        loss_g = tf.reduce_sum(losses_g)\n",
    "        loss_d = tf.reduce_sum(losses_d)\n",
    "        \n",
    "        # update weight of generators\n",
    "        g_grads = g_tape.gradient(losses_g,\n",
    "                                  gen_trainable_variables + noise_generator.trainable_variables)\n",
    "        \n",
    "        StackGANGeneratorWrapper.optimizer\\\n",
    "            .apply_gradients(zip(g_grads, gen_trainable_variables + noise_generator.trainable_variables))\n",
    "\n",
    "        \n",
    "        # update weight of discriminators\n",
    "        d_grads = d_tape.gradient(losses_d, discr_trainable_variables)\n",
    "        \n",
    "        StackGANDiscriminatorWrapper.optimizer\\\n",
    "            .apply_gradients(zip(d_grads, discr_trainable_variables))\n",
    "\n",
    "        return losses_g, losses_d\n",
    "        \n",
    "    return _train_step_template"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "axR-vkyz6kZa",
    "pycharm": {}
   },
   "source": [
    "Inference step\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "r57i63lv6mRT",
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "def stack_gan_gen_step_template(generators, noise_gen, settings):\n",
    "    \n",
    "    is_conditional = settings.is_conditional\n",
    "    branch_num = settings.branch_num\n",
    "    training = False\n",
    "    @tf.function\n",
    "    def _gen_step_template(images, embeddings):\n",
    "        imgs = []\n",
    "        \n",
    "        gen_64_input = tf.concat([embeddings, tf.shape(noise_gen(images))[0]], axis=1)\n",
    "        gen_64_out, imgs_64 = generators['64'].model(gen_64_input, training=training)\n",
    "        imgs.append(imgs_64)\n",
    "        \n",
    "        if branch_num > 1:\n",
    "            gen_128_input = GeneratorNextStage.concat_embedding_and_previous_output(\n",
    "                        embedding=embeddings,\n",
    "                        previous_output=gen_64_out)\n",
    "            gen_128_out, imgs_128 = generators['128'].model(gen_128_input, training=training)\n",
    "            imgs.append(imgs_128)\n",
    "        \n",
    "        if branch_num > 2:\n",
    "            gen_256_input = GeneratorNextStage.concat_embedding_and_previous_output(\n",
    "                       embedding=embeddings,\n",
    "                       previous_output=gen_128_out)\n",
    "            _, imgs_256 = generators['256'].model(gen_256_input, training=training)\n",
    "            imgs.append(imgs_256)\n",
    "        \n",
    "        return [convert_from_output_to_image(img_batch) for img_batch in imgs]\n",
    "\n",
    "    return _gen_step_template"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8xRLPJoW7IcI",
    "pycharm": {}
   },
   "source": [
    "Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ut1yWYzD7JeR",
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "def train(train_step, gen_step, epochs, data, settings, ckptManager, tensorboard_manager, save_images=True, display_images=True):\n",
    "    imgs_sizes = [x[0] for x in zip(['64', '128', '256'], range(settings.branch_num))]\n",
    "    print('img sizes to generate', imgs_sizes)\n",
    "    generator_err_descr = ['generator_{}'.format(x) for x in imgs_sizes]\n",
    "    discriminator_err_descr = ['discriminator_{}'.format(x) for x in imgs_sizes]\n",
    "    timestamp = get_time()\n",
    "    max_batch_iter = 0\n",
    "    checkpoint_save_iteration = 5000 // settings.batch_size # save checkpoint every 5k images\n",
    "    tensorboard_save_iteration = 1000 // settings.batch_size # save tensorboard logs every 1k images\n",
    "    epoch_index = read_epoch_index(settings)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        batch_iter = 0\n",
    "        epoch_start = time.time()\n",
    "        for batch in data:\n",
    "            images, embedding = prepare_embeddings_from_batch(batch,\n",
    "                                                         attributes_subset=attributes_subset)\n",
    "            batch_iter+=1\n",
    "            generator_err, discriminator_err = train_step(images, embedding)\n",
    "            \n",
    "            if (batch_iter+1)%tensorboard_save_iteration == 0:\n",
    "                tensorboard_manager.save('train',\n",
    "                                         [np.average(loss) for loss in generator_err],\n",
    "                                         generator_err_descr,\n",
    "                                         timestamp,\n",
    "                                         'scalars',\n",
    "                                         epoch*max_batch_iter+batch_iter\n",
    "                                        )\n",
    "                tensorboard_manager.save('train',\n",
    "                                         [np.average(loss) for loss in discriminator_err],\n",
    "                                         discriminator_err_descr,\n",
    "                                         timestamp,\n",
    "                                         'scalars',\n",
    "                                         epoch*max_batch_iter+batch_iter\n",
    "                                        )\n",
    "            if (batch_iter+1) % checkpoint_save_iteration == 0:\n",
    "                print('generator loss', np.sum(generator_err), 'discriminator loss', np.sum(discriminator_err))\n",
    "                print(\"Saving checkpoint, iter\", batch_iter+1)\n",
    "                ckptManager.save()\n",
    "        \n",
    "        # save the number of batches in one epoch\n",
    "        if epoch == 0:\n",
    "            max_batch_iter = batch_iter\n",
    "\n",
    "        epoch_end = time.time()\n",
    "        print('-'*30)\n",
    "        save_epoch_index(settings, epoch_index+epoch+1)\n",
    "        print('Epoch {0}/{1}, duration {2}'.format(epoch+1, epochs, epoch_end-epoch_start))\n",
    "        print(\"Saving checkpoint, epoch\", epoch+1)\n",
    "        ckptManager.save()\n",
    "\n",
    "        imgs, embeddings = [prepare_embeddings_from_batch(batch,\n",
    "                                                         attributes_subset=attributes_subset) for batch in data.take(1)][0] # take one batch from train_data\n",
    "        generated_imgs = gen_step(imgs, embeddings)\n",
    "        \n",
    "        for (generated, size) in zip(generated_imgs, imgs_sizes):\n",
    "            tensorboard_manager.save('train',\n",
    "                                     generated,\n",
    "                                     '{}_{}_{}/{}'.format(settings.run_name, size, epoch+1, epochs),\n",
    "                                     timestamp,\n",
    "                                     'images',\n",
    "                                     epoch\n",
    "                                    )\n",
    "            \n",
    "\n",
    "            show_images(generated, epoch, settings, save_images=save_images, display_images=display_images)\n",
    "        print('+'*30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i1EZ3WMg7ZS0",
    "pycharm": {}
   },
   "source": [
    "Training with Wasserstein loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "class StackGANWrapper:\n",
    "    def __init__(self, settings):\n",
    "        self.blocks = dict()\n",
    "        self.name = ''\n",
    "        self.settings = settings\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        return self.blocks[key]\n",
    "\n",
    "class StackGANGeneratorWrapper(StackGANWrapper):\n",
    "    optimizer = tf.keras.optimizers.Adam(1e-4)\n",
    "    def __init__(self, settings):\n",
    "        super().__init__(settings)\n",
    "        self.name = 'generators'\n",
    "        self.blocks['64'] = ModelWrapper(\n",
    "                        model_object = GeneratorInitStage(noise_shape=settings.z_dim,\n",
    "                            embeddings_shape=settings.embeddings_shape,\n",
    "                            filters=settings.gen_init_filters,\n",
    "                            is_conditional=settings.is_conditional,\n",
    "                            settings=settings,\n",
    "                            model_name='gen_64'),\n",
    "                        model_optimizer_fun = StackGANGeneratorWrapper.optimizer,\n",
    "                        model_loss_fun = min_max_generator_loss\n",
    "                        )\n",
    "        \n",
    "        self.blocks['128'] = ModelWrapper(\n",
    "                        model_object = GeneratorNextStage(\n",
    "                            embeddings_shape=settings.embeddings_shape,\n",
    "                            number_of_residuals=2,\n",
    "                            previous_stage_shape=self['64'].model.output_shape,\n",
    "                            filters=settings.gf_dim,\n",
    "                            is_conditional=settings.is_conditional,\n",
    "                            model_name='gen_128'),\n",
    "                        model_optimizer_fun = StackGANGeneratorWrapper.optimizer,\n",
    "                        model_loss_fun = min_max_generator_loss\n",
    "                        )\n",
    "                    \n",
    "        self.blocks['256'] = ModelWrapper(\n",
    "                        model_object = GeneratorNextStage(\n",
    "                            embeddings_shape=settings.embeddings_shape,\n",
    "                            number_of_residuals=2,\n",
    "                            previous_stage_shape=self['128'].model.output_shape,\n",
    "                            filters=settings.gf_dim,\n",
    "                            is_conditional=settings.is_conditional,\n",
    "                            model_name='gen_256'),\n",
    "                        model_optimizer_fun = StackGANGeneratorWrapper.optimizer,\n",
    "                        model_loss_fun = min_max_generator_loss\n",
    "                        )\n",
    "        \n",
    "class StackGANDiscriminatorWrapper(StackGANWrapper):\n",
    "    optimizer = tf.keras.optimizers.Adam(1e-4)\n",
    "    def __init__(self, generators, settings):\n",
    "        super().__init__(settings)\n",
    "        self.name = 'discriminators'\n",
    "        self.blocks['64'] = ModelWrapper(\n",
    "            model_object = StackGANDiscriminator(\n",
    "                input_shape=generators['64'].model.image_shape,\n",
    "                filters=settings.df_dim,\n",
    "                embeddings_shape=settings.embeddings_shape,\n",
    "                output_image_size=64,\n",
    "                is_conditional=settings.is_conditional,\n",
    "                model_name='discr_64'\n",
    "            ),\n",
    "            model_optimizer_fun = StackGANDiscriminatorWrapper.optimizer,\n",
    "            model_loss_fun = stack_gan_discriminator_loss\n",
    "        )\n",
    "\n",
    "        self.blocks['128'] = ModelWrapper(\n",
    "            model_object = StackGANDiscriminator(\n",
    "                input_shape=generators['128'].model.image_shape,\n",
    "                filters=settings.df_dim,\n",
    "                embeddings_shape=settings.embeddings_shape,\n",
    "                output_image_size=128,\n",
    "                is_conditional=settings.is_conditional,\n",
    "                model_name='discr_128'\n",
    "            ),\n",
    "            model_optimizer_fun = StackGANDiscriminatorWrapper.optimizer,\n",
    "            model_loss_fun = stack_gan_discriminator_loss\n",
    "        )\n",
    "        \n",
    "        self.blocks['256'] = ModelWrapper(\n",
    "            model_object = StackGANDiscriminator(\n",
    "                input_shape=generators['256'].model.image_shape,\n",
    "                filters=settings.df_dim,\n",
    "                embeddings_shape=settings.embeddings_shape,\n",
    "                output_image_size=256,\n",
    "                is_conditional=settings.is_conditional,\n",
    "                model_name='discr_256'\n",
    "            ),\n",
    "            model_optimizer_fun = StackGANDiscriminatorWrapper.optimizer,\n",
    "            model_loss_fun = stack_gan_discriminator_loss\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "EHKejxSy7YOx",
    "outputId": "de999bc7-5064-4d75-86e6-487f676dce28",
    "pycharm": {},
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "env = Environment(collab_mode)\n",
    "generators = StackGANGeneratorWrapper(env.settings)\n",
    "discriminators = StackGANDiscriminatorWrapper(generators, env.settings)\n",
    "env.settings.epochs = 25\n",
    "env.settings.branch_num = 1\n",
    "env.settings.subdataset_dir='img_align_celeba'\n",
    "env.settings.save_models = True\n",
    "print(env.settings.get_base_path)\n",
    "# load data\n",
    "env.datasetCache.load_data(env.settings)\n",
    "# set models\n",
    "env.models['generators'] = generators\n",
    "env.models['discriminators'] = discriminators\n",
    "# setup tensorboard\n",
    "env.tensorboard = TensorboardManager()\n",
    "env.tensorboard.initialize(env.settings)\n",
    "\n",
    "noise = StackGANNoiseGenerator(100)\n",
    "\n",
    "if env.settings.branch_num == 1:\n",
    "    checkpoint = tf.train.Checkpoint(generator_opt=StackGANGeneratorWrapper.optimizer,\n",
    "                                         discriminator_opt=StackGANDiscriminatorWrapper.optimizer,\n",
    "                                         generator_64=env.models['generators']['64'].model,\n",
    "                                         discriminator_64=env.models['discriminators']['64'].model)\n",
    "if env.settings.branch_num == 2:\n",
    "    checkpoint = tf.train.Checkpoint(generator_opt=StackGANGeneratorWrapper.optimizer,\n",
    "                                     discriminator_opt=StackGANDiscriminatorWrapper.optimizer,\n",
    "                                     generator_64=env.models['generators']['64'].model,\n",
    "                                     generator_128=env.models['generators']['128'].model,\n",
    "                                     discriminator_64=env.models['discriminators']['64'].model,\n",
    "                                     discriminator_128=env.models['discriminators']['128'].model)\n",
    "if env.settings.branch_num == 3:\n",
    "    checkpoint = tf.train.Checkpoint(generator_opt=StackGANGeneratorWrapper.optimizer,\n",
    "                                     discriminator_opt=StackGANDiscriminatorWrapper.optimizer,\n",
    "                                     generator_64=env.models['generators']['64'].model,\n",
    "                                     generator_128=env.models['generators']['128'].model,\n",
    "                                     generator_256=env.models['generators']['256'].model,\n",
    "                                     discriminator_64=env.models['discriminators']['64'].model,\n",
    "                                     discriminator_128=env.models['discriminators']['128'].model,\n",
    "                                     discriminator_256=env.models['discriminators']['256'].model)\n",
    "\n",
    "env.checkpointManager = tf.train.CheckpointManager(checkpoint=checkpoint,\n",
    "                                                   directory=env.settings.checkpoint_dir,\n",
    "                                                   max_to_keep=3\n",
    "                                                  )\n",
    "if env.checkpointManager.latest_checkpoint:\n",
    "    print(\"restoring state from\", env.checkpointManager.latest_checkpoint)\n",
    "    checkpoint\\\n",
    "        .restore(env.checkpointManager.latest_checkpoint)\n",
    "else:\n",
    "    print('not restored')\n",
    "    \n",
    "train_step = stack_gan_train_step_template(\n",
    "    generators=env.models['generators'],\n",
    "    discriminators=env.models['discriminators'],\n",
    "    noise_generator=noise,\n",
    "    settings=env.settings\n",
    ")\n",
    "\n",
    "gen_step = stack_gan_gen_step_template(\n",
    "    generators=env.models['generators'],\n",
    "    noise_gen=noise,\n",
    "    settings=env.settings\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "print(\"Start time\", get_time())\n",
    "print('%'*30)\n",
    "start = time.time()\n",
    "\n",
    "train(\n",
    "    train_step=train_step,\n",
    "    gen_step=gen_step,\n",
    "    epochs=env.settings.epochs,\n",
    "    data=env.datasetCache.data,\n",
    "    settings=env.settings,\n",
    "    ckptManager=env.checkpointManager,\n",
    "    tensorboard_manager=env.tensorboard\n",
    ")\n",
    "\n",
    "end = time.time()\n",
    "print('%'*30)\n",
    "print(\"End time\", get_time())\n",
    "print(\"seconds elapsed\", end - start)\n",
    "\n",
    "if env.settings.save_models:\n",
    "    print('saving models')\n",
    "    for model in env.models:\n",
    "        if isinstance(env.models[model], StackGANWrapper):\n",
    "            # models are wrapped in StackGANWrapper\n",
    "            for model_ in env.models[model].blocks.values():\n",
    "                model_.model.save_model(env.settings.model_save_path)\n",
    "        elif isinstance(env.models[model], ModelWrapper):\n",
    "            # models are wrapped in ModelWrapper\n",
    "            env.models[model].model.save_model(env.settings.model_save_path)\n",
    "        else:\n",
    "            # models are stored explicitly\n",
    "            env.models[model].save_model(env.settings.model_save_path)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "FaceGenerator",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "FaceGenerator",
   "language": "python",
   "name": "facegenerator"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}