"""Collection of generative models."""

import torch as th
import ttools

import rendering
import modules

LOG = ttools.get_logger(__name__)


class BaseModel(th.nn.Module):
    def sample_z(self, bs, device="cpu"):
        return th.randn(bs, self.zdim).to(device)


class BaseVectorModel(BaseModel):
    def get_vector(self, z):
        _, scenes = self._forward(z)
        return scenes

    def _forward(self, x):
        raise NotImplementedError()

    def forward(self, z):
        # Only return the raster
        return self._forward(z)[0]


class BezierVectorGenerator(BaseVectorModel):
    NUM_SEGMENTS = 2
    def __init__(self, num_strokes=4,
                 zdim=128, width=32, imsize=32,
                 color_output=False,
                 stroke_width=None):
        super(BezierVectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim

        self.trunk = th.nn.Sequential(
            th.nn.Linear(zdim, width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(width, 2*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(2*width, 4*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(4*width, 8*width),
            th.nn.SELU(inplace=True),
        )

        # 4 points bezier with n_segments -> 3*n_segments + 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, 
                         2*self.num_strokes*(
                             BezierVectorGenerator.NUM_SEGMENTS*3 + 1)),
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(
                th.nn.Linear(8*width, 3*self.num_strokes),
                th.nn.Sigmoid()
            )

    def _forward(self, z):
        bs = z.shape[0]

        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)

        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None

        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        all_points = all_points.view(
            bs, self.num_strokes, BezierVectorGenerator.NUM_SEGMENTS*3+1, 2)

        output, scenes = rendering.bezier_render(all_points, all_widths, all_alphas,
                                         colors=all_colors,
                                         canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class VectorGenerator(BaseVectorModel):
    def __init__(self, num_strokes=4,
                 zdim=128, width=32, imsize=32,
                 color_output=False,
                 stroke_width=None):
        super(VectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim

        self.trunk = th.nn.Sequential(
            th.nn.Linear(zdim, width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(width, 2*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(2*width, 4*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(4*width, 8*width),
            th.nn.SELU(inplace=True),
        )

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, 2*(self.num_strokes*2)),
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(
                th.nn.Linear(8*width, 3*self.num_strokes),
                th.nn.Sigmoid()
            )

    def _forward(self, z):
        bs = z.shape[0]

        feats = self.trunk(z)

        all_points = self.point_predictor(feats)

        all_alphas = self.alpha_predictor(feats)

        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None

        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        all_points = all_points.view(bs, self.num_strokes, 2, 2)
        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                       colors=all_colors,
                                       canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class RNNVectorGenerator(BaseVectorModel):
    def __init__(self, num_strokes=64,
                 zdim=128, width=32, imsize=32,
                 hidden_size=512, dropout=0.9,
                 color_output=False,
                 num_layers=3, stroke_width=None):
        super(RNNVectorGenerator, self).__init__()


        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim

        self.hidden_cell_predictor = th.nn.Linear(
            zdim, 2*hidden_size*num_layers)

        self.lstm = th.nn.LSTM(
            zdim, hidden_size,
            num_layers=self.num_layers, dropout=dropout,
            batch_first=True)

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 2*2),  # 2 points, (x,y)
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes

        # z is passed at each step, duplicate it
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)

        # First step in the RNN
        if hidden_and_cell is None:
            # Initialize from latent vector
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size*self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size*self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = (hidden, cell)

        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell

        feats = feats.reshape(bs*steps, self.hidden_size)

        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)

        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                        canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""
    def __init__(self, num_strokes=64,
                 zdim=128, width=32, imsize=32,
                 hidden_size=512, dropout=0.9,
                 color_output=False,
                 num_layers=3, stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim

        self.hidden_cell_predictor = th.nn.Linear(
            zdim, 2*hidden_size*num_layers)

        self.lstm = th.nn.LSTM(
            zdim, hidden_size,
            num_layers=self.num_layers, dropout=dropout,
            batch_first=True)

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 2),  # 1 point, (x,y)
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes

        # z is passed at each step, duplicate it
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)

        # First step in the RNN
        if hidden_and_cell is None:
            # Initialize from latent vector
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size*self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size*self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = (hidden, cell)

        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell

        feats = feats.reshape(bs*steps, self.hidden_size)

        # Construct the chain
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([
            # first point is canvas center
            th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)

        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)

        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                        canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class Generator(BaseModel):
    def __init__(self, width=64, imsize=32, zdim=128,
                 stroke_width=None,
                 color_output=False,
                 num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32

        self.imsize = imsize
        self.zdim = zdim

        num_in_chans = self.zdim // (2*2)
        num_out_chans = 3 if color_output else 1

        self.net = th.nn.Sequential(
            th.nn.ConvTranspose2d(num_in_chans, width*8, 4, padding=1,
                                  stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width*8, width*8, 3, padding=1),
            th.nn.BatchNorm2d(width*8),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 4x4

            th.nn.ConvTranspose2d(8*width, 4*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(4*width, 4*width, 3, padding=1),
            th.nn.BatchNorm2d(width*4),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 8x8

            th.nn.ConvTranspose2d(4*width, 2*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(2*width, 2*width, 3, padding=1),
            th.nn.BatchNorm2d(width*2),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 16x16

            th.nn.ConvTranspose2d(2*width, width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 32x32

            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, num_out_chans, 1),

            th.nn.Tanh(),
        )

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2*2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):
    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()

        self.conditional = conditional

        sn = th.nn.utils.spectral_norm

        num_chan_in = 3 if color_output else 1

        self.net = th.nn.Sequential(
            th.nn.Conv2d(num_chan_in, width, 3, padding=1),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, 2*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 16x16

            sn(th.nn.Conv2d(2*width, 2*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(2*width, 4*width, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 8x8

            sn(th.nn.Conv2d(4*width, 4*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(4*width, width*4, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 4x4

            sn(th.nn.Conv2d(4*width, 4*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(4*width, width*4, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 2x2

            modules.Flatten(),
            th.nn.Linear(width*4*2*2, 1),
        )

    def forward(self, x):
        out = self.net(x)
        return out

0267a290a67a3960c700d317c5a832770d756753
<mask token>


class VectorGenerator(BaseVectorModel):
    <mask token>
    <mask token>


class RNNVectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(RNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2 *
            2), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class Generator(BaseModel):

    def __init__(self, width=64, imsize=32, zdim=128, stroke_width=None,
        color_output=False, num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32
        self.imsize = imsize
        self.zdim = zdim
        num_in_chans = self.zdim // (2 * 2)
        num_out_chans = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.ConvTranspose2d(num_in_chans, 
            width * 8, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width * 8, width * 8, 3, padding=1),
            th.nn.BatchNorm2d(width * 8), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(8 * width, 4 * width, 4, padding=1,
            stride=2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(4 *
            width, 4 * width, 3, padding=1), th.nn.BatchNorm2d(width * 4),
            th.nn.LeakyReLU(0.2, inplace=True), th.nn.ConvTranspose2d(4 *
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(2 * width, 2 * width, 3, padding=1),
            th.nn.BatchNorm2d(width * 2), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(2 * width, width, 4, padding=1, stride
            =2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(width,
            width, 3, padding=1), th.nn.BatchNorm2d(width), th.nn.LeakyReLU
            (0.2, inplace=True), th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width), th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width, num_out_chans, 1), th.nn.Tanh())

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2 * 2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):

    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()
        self.conditional = conditional
        sn = th.nn.utils.spectral_norm
        num_chan_in = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.Conv2d(num_chan_in, width, 3,
            padding=1), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(2 * width, 2 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(2 *
            width, 4 * width, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), modules.Flatten(), th.nn.Linear(width * 4 * 2 * 
            2, 1))

    def forward(self, x):
        out = self.net(x)
        return out

<mask token>


class BaseVectorModel(BaseModel):
    <mask token>

    def _forward(self, x):
        raise NotImplementedError()

    def forward(self, z):
        return self._forward(z)[0]


class BezierVectorGenerator(BaseVectorModel):
    NUM_SEGMENTS = 2

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(BezierVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            self.num_strokes * (BezierVectorGenerator.NUM_SEGMENTS * 3 + 1)
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 
            BezierVectorGenerator.NUM_SEGMENTS * 3 + 1, 2)
        output, scenes = rendering.bezier_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class VectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(VectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            (self.num_strokes * 2)), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 2, 2)
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class RNNVectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(RNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2 *
            2), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class Generator(BaseModel):

    def __init__(self, width=64, imsize=32, zdim=128, stroke_width=None,
        color_output=False, num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32
        self.imsize = imsize
        self.zdim = zdim
        num_in_chans = self.zdim // (2 * 2)
        num_out_chans = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.ConvTranspose2d(num_in_chans, 
            width * 8, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width * 8, width * 8, 3, padding=1),
            th.nn.BatchNorm2d(width * 8), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(8 * width, 4 * width, 4, padding=1,
            stride=2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(4 *
            width, 4 * width, 3, padding=1), th.nn.BatchNorm2d(width * 4),
            th.nn.LeakyReLU(0.2, inplace=True), th.nn.ConvTranspose2d(4 *
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(2 * width, 2 * width, 3, padding=1),
            th.nn.BatchNorm2d(width * 2), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(2 * width, width, 4, padding=1, stride
            =2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(width,
            width, 3, padding=1), th.nn.BatchNorm2d(width), th.nn.LeakyReLU
            (0.2, inplace=True), th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width), th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width, num_out_chans, 1), th.nn.Tanh())

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2 * 2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):

    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()
        self.conditional = conditional
        sn = th.nn.utils.spectral_norm
        num_chan_in = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.Conv2d(num_chan_in, width, 3,
            padding=1), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(2 * width, 2 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(2 *
            width, 4 * width, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), modules.Flatten(), th.nn.Linear(width * 4 * 2 * 
            2, 1))

    def forward(self, x):
        out = self.net(x)
        return out

<mask token>


class BaseVectorModel(BaseModel):

    def get_vector(self, z):
        _, scenes = self._forward(z)
        return scenes

    def _forward(self, x):
        raise NotImplementedError()

    def forward(self, z):
        return self._forward(z)[0]


class BezierVectorGenerator(BaseVectorModel):
    NUM_SEGMENTS = 2

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(BezierVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            self.num_strokes * (BezierVectorGenerator.NUM_SEGMENTS * 3 + 1)
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 
            BezierVectorGenerator.NUM_SEGMENTS * 3 + 1, 2)
        output, scenes = rendering.bezier_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class VectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(VectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            (self.num_strokes * 2)), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 2, 2)
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class RNNVectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(RNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2 *
            2), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class Generator(BaseModel):

    def __init__(self, width=64, imsize=32, zdim=128, stroke_width=None,
        color_output=False, num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32
        self.imsize = imsize
        self.zdim = zdim
        num_in_chans = self.zdim // (2 * 2)
        num_out_chans = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.ConvTranspose2d(num_in_chans, 
            width * 8, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width * 8, width * 8, 3, padding=1),
            th.nn.BatchNorm2d(width * 8), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(8 * width, 4 * width, 4, padding=1,
            stride=2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(4 *
            width, 4 * width, 3, padding=1), th.nn.BatchNorm2d(width * 4),
            th.nn.LeakyReLU(0.2, inplace=True), th.nn.ConvTranspose2d(4 *
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(2 * width, 2 * width, 3, padding=1),
            th.nn.BatchNorm2d(width * 2), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(2 * width, width, 4, padding=1, stride
            =2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(width,
            width, 3, padding=1), th.nn.BatchNorm2d(width), th.nn.LeakyReLU
            (0.2, inplace=True), th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width), th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width, num_out_chans, 1), th.nn.Tanh())

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2 * 2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):

    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()
        self.conditional = conditional
        sn = th.nn.utils.spectral_norm
        num_chan_in = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.Conv2d(num_chan_in, width, 3,
            padding=1), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(2 * width, 2 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(2 *
            width, 4 * width, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), modules.Flatten(), th.nn.Linear(width * 4 * 2 * 
            2, 1))

    def forward(self, x):
        out = self.net(x)
        return out

<mask token>
import torch as th
import ttools
import rendering
import modules
LOG = ttools.get_logger(__name__)


class BaseModel(th.nn.Module):

    def sample_z(self, bs, device='cpu'):
        return th.randn(bs, self.zdim).to(device)


class BaseVectorModel(BaseModel):

    def get_vector(self, z):
        _, scenes = self._forward(z)
        return scenes

    def _forward(self, x):
        raise NotImplementedError()

    def forward(self, z):
        return self._forward(z)[0]


class BezierVectorGenerator(BaseVectorModel):
    NUM_SEGMENTS = 2

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(BezierVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            self.num_strokes * (BezierVectorGenerator.NUM_SEGMENTS * 3 + 1)
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 
            BezierVectorGenerator.NUM_SEGMENTS * 3 + 1, 2)
        output, scenes = rendering.bezier_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class VectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=4, zdim=128, width=32, imsize=32,
        color_output=False, stroke_width=None):
        super(VectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim
        self.trunk = th.nn.Sequential(th.nn.Linear(zdim, width), th.nn.SELU
            (inplace=True), th.nn.Linear(width, 2 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(2 * width, 4 * width), th.nn.SELU(
            inplace=True), th.nn.Linear(4 * width, 8 * width), th.nn.SELU(
            inplace=True))
        self.point_predictor = th.nn.Sequential(th.nn.Linear(8 * width, 2 *
            (self.num_strokes * 2)), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
            self.num_strokes), th.nn.Sigmoid())
        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(th.nn.Linear(8 * width,
                3 * self.num_strokes), th.nn.Sigmoid())

    def _forward(self, z):
        bs = z.shape[0]
        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)
        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None
        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        all_points = all_points.view(bs, self.num_strokes, 2, 2)
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, colors=all_colors, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class RNNVectorGenerator(BaseVectorModel):

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(RNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2 *
            2), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""

    def __init__(self, num_strokes=64, zdim=128, width=32, imsize=32,
        hidden_size=512, dropout=0.9, color_output=False, num_layers=3,
        stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()
        if stroke_width is None:
            self.stroke_width = 0.5, 3.0
            LOG.warning('Setting default stroke with %s', self.stroke_width)
        else:
            self.stroke_width = stroke_width
        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim
        self.hidden_cell_predictor = th.nn.Linear(zdim, 2 * hidden_size *
            num_layers)
        self.lstm = th.nn.LSTM(zdim, hidden_size, num_layers=self.
            num_layers, dropout=dropout, batch_first=True)
        self.point_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 2
            ), th.nn.Tanh())
        self.width_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())
        self.alpha_predictor = th.nn.Sequential(th.nn.Linear(hidden_size, 1
            ), th.nn.Sigmoid())

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
        if hidden_and_cell is None:
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = hidden, cell
        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell
        feats = feats.reshape(bs * steps, self.hidden_size)
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width
        output, scenes = rendering.line_render(all_points, all_widths,
            all_alphas, canvas_size=self.imsize)
        output = output * 2.0 - 1.0
        return output, scenes


class Generator(BaseModel):

    def __init__(self, width=64, imsize=32, zdim=128, stroke_width=None,
        color_output=False, num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32
        self.imsize = imsize
        self.zdim = zdim
        num_in_chans = self.zdim // (2 * 2)
        num_out_chans = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.ConvTranspose2d(num_in_chans, 
            width * 8, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width * 8, width * 8, 3, padding=1),
            th.nn.BatchNorm2d(width * 8), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(8 * width, 4 * width, 4, padding=1,
            stride=2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(4 *
            width, 4 * width, 3, padding=1), th.nn.BatchNorm2d(width * 4),
            th.nn.LeakyReLU(0.2, inplace=True), th.nn.ConvTranspose2d(4 *
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(2 * width, 2 * width, 3, padding=1),
            th.nn.BatchNorm2d(width * 2), th.nn.LeakyReLU(0.2, inplace=True
            ), th.nn.ConvTranspose2d(2 * width, width, 4, padding=1, stride
            =2), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(width,
            width, 3, padding=1), th.nn.BatchNorm2d(width), th.nn.LeakyReLU
            (0.2, inplace=True), th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width), th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1), th.nn.LeakyReLU(0.2,
            inplace=True), th.nn.Conv2d(width, num_out_chans, 1), th.nn.Tanh())

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2 * 2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):

    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()
        self.conditional = conditional
        sn = th.nn.utils.spectral_norm
        num_chan_in = 3 if color_output else 1
        self.net = th.nn.Sequential(th.nn.Conv2d(num_chan_in, width, 3,
            padding=1), th.nn.LeakyReLU(0.2, inplace=True), th.nn.Conv2d(
            width, 2 * width, 4, padding=1, stride=2), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(2 * width, 2 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(2 *
            width, 4 * width, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), sn(th.nn.Conv2d(4 * width, 4 * width, 3, padding
            =1)), th.nn.LeakyReLU(0.2, inplace=True), sn(th.nn.Conv2d(4 *
            width, width * 4, 4, padding=1, stride=2)), th.nn.LeakyReLU(0.2,
            inplace=True), modules.Flatten(), th.nn.Linear(width * 4 * 2 * 
            2, 1))

    def forward(self, x):
        out = self.net(x)
        return out

"""Collection of generative models."""

import torch as th
import ttools

import rendering
import modules

LOG = ttools.get_logger(__name__)


class BaseModel(th.nn.Module):
    def sample_z(self, bs, device="cpu"):
        return th.randn(bs, self.zdim).to(device)


class BaseVectorModel(BaseModel):
    def get_vector(self, z):
        _, scenes = self._forward(z)
        return scenes

    def _forward(self, x):
        raise NotImplementedError()

    def forward(self, z):
        # Only return the raster
        return self._forward(z)[0]


class BezierVectorGenerator(BaseVectorModel):
    NUM_SEGMENTS = 2
    def __init__(self, num_strokes=4,
                 zdim=128, width=32, imsize=32,
                 color_output=False,
                 stroke_width=None):
        super(BezierVectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim

        self.trunk = th.nn.Sequential(
            th.nn.Linear(zdim, width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(width, 2*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(2*width, 4*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(4*width, 8*width),
            th.nn.SELU(inplace=True),
        )

        # 4 points bezier with n_segments -> 3*n_segments + 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, 
                         2*self.num_strokes*(
                             BezierVectorGenerator.NUM_SEGMENTS*3 + 1)),
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(
                th.nn.Linear(8*width, 3*self.num_strokes),
                th.nn.Sigmoid()
            )

    def _forward(self, z):
        bs = z.shape[0]

        feats = self.trunk(z)
        all_points = self.point_predictor(feats)
        all_alphas = self.alpha_predictor(feats)

        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None

        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        all_points = all_points.view(
            bs, self.num_strokes, BezierVectorGenerator.NUM_SEGMENTS*3+1, 2)

        output, scenes = rendering.bezier_render(all_points, all_widths, all_alphas,
                                         colors=all_colors,
                                         canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class VectorGenerator(BaseVectorModel):
    def __init__(self, num_strokes=4,
                 zdim=128, width=32, imsize=32,
                 color_output=False,
                 stroke_width=None):
        super(VectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.imsize = imsize
        self.num_strokes = num_strokes
        self.zdim = zdim

        self.trunk = th.nn.Sequential(
            th.nn.Linear(zdim, width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(width, 2*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(2*width, 4*width),
            th.nn.SELU(inplace=True),

            th.nn.Linear(4*width, 8*width),
            th.nn.SELU(inplace=True),
        )

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, 2*(self.num_strokes*2)),
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(8*width, self.num_strokes),
            th.nn.Sigmoid()
        )

        self.color_predictor = None
        if color_output:
            self.color_predictor = th.nn.Sequential(
                th.nn.Linear(8*width, 3*self.num_strokes),
                th.nn.Sigmoid()
            )

    def _forward(self, z):
        bs = z.shape[0]

        feats = self.trunk(z)

        all_points = self.point_predictor(feats)

        all_alphas = self.alpha_predictor(feats)

        if self.color_predictor:
            all_colors = self.color_predictor(feats)
            all_colors = all_colors.view(bs, self.num_strokes, 3)
        else:
            all_colors = None

        all_widths = self.width_predictor(feats)
        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        all_points = all_points.view(bs, self.num_strokes, 2, 2)
        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                       colors=all_colors,
                                       canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class RNNVectorGenerator(BaseVectorModel):
    def __init__(self, num_strokes=64,
                 zdim=128, width=32, imsize=32,
                 hidden_size=512, dropout=0.9,
                 color_output=False,
                 num_layers=3, stroke_width=None):
        super(RNNVectorGenerator, self).__init__()


        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim

        self.hidden_cell_predictor = th.nn.Linear(
            zdim, 2*hidden_size*num_layers)

        self.lstm = th.nn.LSTM(
            zdim, hidden_size,
            num_layers=self.num_layers, dropout=dropout,
            batch_first=True)

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 2*2),  # 2 points, (x,y)
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes

        # z is passed at each step, duplicate it
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)

        # First step in the RNN
        if hidden_and_cell is None:
            # Initialize from latent vector
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size*self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size*self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = (hidden, cell)

        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell

        feats = feats.reshape(bs*steps, self.hidden_size)

        all_points = self.point_predictor(feats).view(bs, steps, 2, 2)
        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)

        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                        canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class ChainRNNVectorGenerator(BaseVectorModel):
    """Strokes form a single long chain."""
    def __init__(self, num_strokes=64,
                 zdim=128, width=32, imsize=32,
                 hidden_size=512, dropout=0.9,
                 color_output=False,
                 num_layers=3, stroke_width=None):
        super(ChainRNNVectorGenerator, self).__init__()

        if stroke_width is None:
            self.stroke_width = (0.5, 3.0)
            LOG.warning("Setting default stroke with %s", self.stroke_width)
        else:
            self.stroke_width = stroke_width

        self.num_layers = num_layers
        self.imsize = imsize
        self.num_strokes = num_strokes
        self.hidden_size = hidden_size
        self.zdim = zdim

        self.hidden_cell_predictor = th.nn.Linear(
            zdim, 2*hidden_size*num_layers)

        self.lstm = th.nn.LSTM(
            zdim, hidden_size,
            num_layers=self.num_layers, dropout=dropout,
            batch_first=True)

        # straight lines so n_segments -> n_segments - 1 points
        self.point_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 2),  # 1 point, (x,y)
            th.nn.Tanh()  # bound spatial extent
        )

        self.width_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

        self.alpha_predictor = th.nn.Sequential(
            th.nn.Linear(hidden_size, 1),
            th.nn.Sigmoid()
        )

    def _forward(self, z, hidden_and_cell=None):
        steps = self.num_strokes

        # z is passed at each step, duplicate it
        bs = z.shape[0]
        expanded_z = z.unsqueeze(1).repeat(1, steps, 1)

        # First step in the RNN
        if hidden_and_cell is None:
            # Initialize from latent vector
            hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
            hidden = hidden_and_cell[:, :self.hidden_size*self.num_layers]
            hidden = hidden.view(-1, self.num_layers, self.hidden_size)
            hidden = hidden.permute(1, 0, 2).contiguous()
            cell = hidden_and_cell[:, self.hidden_size*self.num_layers:]
            cell = cell.view(-1, self.num_layers, self.hidden_size)
            cell = cell.permute(1, 0, 2).contiguous()
            hidden_and_cell = (hidden, cell)

        feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
        hidden, cell = hidden_and_cell

        feats = feats.reshape(bs*steps, self.hidden_size)

        # Construct the chain
        end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
        start_points = th.cat([
            # first point is canvas center
            th.zeros(bs, 1, 1, 2, device=feats.device),
            end_points[:, 1:, :, :]], 1)
        all_points = th.cat([start_points, end_points], 2)

        all_alphas = self.alpha_predictor(feats).view(bs, steps)
        all_widths = self.width_predictor(feats).view(bs, steps)

        min_width = self.stroke_width[0]
        max_width = self.stroke_width[1]
        all_widths = (max_width - min_width) * all_widths + min_width

        output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
                                        canvas_size=self.imsize)

        # map to [-1, 1]
        output = output*2.0 - 1.0

        return output, scenes


class Generator(BaseModel):
    def __init__(self, width=64, imsize=32, zdim=128,
                 stroke_width=None,
                 color_output=False,
                 num_strokes=4):
        super(Generator, self).__init__()
        assert imsize == 32

        self.imsize = imsize
        self.zdim = zdim

        num_in_chans = self.zdim // (2*2)
        num_out_chans = 3 if color_output else 1

        self.net = th.nn.Sequential(
            th.nn.ConvTranspose2d(num_in_chans, width*8, 4, padding=1,
                                  stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width*8, width*8, 3, padding=1),
            th.nn.BatchNorm2d(width*8),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 4x4

            th.nn.ConvTranspose2d(8*width, 4*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(4*width, 4*width, 3, padding=1),
            th.nn.BatchNorm2d(width*4),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 8x8

            th.nn.ConvTranspose2d(4*width, 2*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(2*width, 2*width, 3, padding=1),
            th.nn.BatchNorm2d(width*2),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 16x16

            th.nn.ConvTranspose2d(2*width, width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 32x32

            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.BatchNorm2d(width),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, width, 3, padding=1),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, num_out_chans, 1),

            th.nn.Tanh(),
        )

    def forward(self, z):
        bs = z.shape[0]
        num_in_chans = self.zdim // (2*2)
        raster = self.net(z.view(bs, num_in_chans, 2, 2))
        return raster


class Discriminator(th.nn.Module):
    def __init__(self, conditional=False, width=64, color_output=False):
        super(Discriminator, self).__init__()

        self.conditional = conditional

        sn = th.nn.utils.spectral_norm

        num_chan_in = 3 if color_output else 1

        self.net = th.nn.Sequential(
            th.nn.Conv2d(num_chan_in, width, 3, padding=1),
            th.nn.LeakyReLU(0.2, inplace=True),
            th.nn.Conv2d(width, 2*width, 4, padding=1, stride=2),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 16x16

            sn(th.nn.Conv2d(2*width, 2*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(2*width, 4*width, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 8x8

            sn(th.nn.Conv2d(4*width, 4*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(4*width, width*4, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 4x4

            sn(th.nn.Conv2d(4*width, 4*width, 3, padding=1)),
            th.nn.LeakyReLU(0.2, inplace=True),
            sn(th.nn.Conv2d(4*width, width*4, 4, padding=1, stride=2)),
            th.nn.LeakyReLU(0.2, inplace=True),
            # 2x2

            modules.Flatten(),
            th.nn.Linear(width*4*2*2, 1),
        )

    def forward(self, x):
        out = self.net(x)
        return out
