Create a :class:`StyleTransfer` model.

    Parameters
    ----------
    style_dataset: SFrame
        Input style images. The columns named by the ``style_feature`` parameters will
        be extracted for training the model.

    content_dataset : SFrame
        Input content images. The columns named by the ``content_feature`` parameters will
        be extracted for training the model.

    style_feature: string
        Name of the column containing the input images in style SFrame.
        'None' (the default) indicates the only image column in the style SFrame
        should be used as the feature.

    content_feature: string
        Name of the column containing the input images in content SFrame.
        'None' (the default) indicates the only image column in the content
        SFrame should be used as the feature.

    max_iterations : int
        The number of training iterations. If 'None' (the default), then it will
        be automatically determined based on the amount of data you provide.

    model : string optional
        Style transfer model to use:

            - "resnet-16" : Fast and small-sized residual network that uses
                            VGG-16 as reference network during training.

    batch_size : int, optional
        If you are getting memory errors, try decreasing this value. If you
        have a powerful computer, increasing this value may improve training
        throughput.

    verbose : bool, optional
        If True, print progress updates and model details.


    Returns
    -------
    out : StyleTransfer
        A trained :class:`StyleTransfer` model.

    See Also
    --------
    StyleTransfer

    Examples
    --------
    .. sourcecode:: python

        # Create datasets
        >>> content_dataset = turicreate.image_analysis.load_images('content_images/')
        >>> style_dataset = turicreate.image_analysis.load_images('style_images/')

        # Train a style transfer model
        >>> model = turicreate.style_transfer.create(content_dataset, style_dataset)

        # Stylize an image on all styles
        >>> stylized_images = model.stylize(data)

        # Visualize the stylized images
        >>> stylized_images.explore()

def create(style_dataset, content_dataset, style_feature=None,
        content_feature=None, max_iterations=None, model='resnet-16',
        verbose=True, batch_size = 6, **kwargs):
    """
    Create a :class:`StyleTransfer` model.

    Parameters
    ----------
    style_dataset: SFrame
        Input style images. The columns named by the ``style_feature`` parameters will
        be extracted for training the model.

    content_dataset : SFrame
        Input content images. The columns named by the ``content_feature`` parameters will
        be extracted for training the model.

    style_feature: string
        Name of the column containing the input images in style SFrame.
        'None' (the default) indicates the only image column in the style SFrame
        should be used as the feature.

    content_feature: string
        Name of the column containing the input images in content SFrame.
        'None' (the default) indicates the only image column in the content
        SFrame should be used as the feature.

    max_iterations : int
        The number of training iterations. If 'None' (the default), then it will
        be automatically determined based on the amount of data you provide.

    model : string optional
        Style transfer model to use:

            - "resnet-16" : Fast and small-sized residual network that uses
                            VGG-16 as reference network during training.

    batch_size : int, optional
        If you are getting memory errors, try decreasing this value. If you
        have a powerful computer, increasing this value may improve training
        throughput.

    verbose : bool, optional
        If True, print progress updates and model details.


    Returns
    -------
    out : StyleTransfer
        A trained :class:`StyleTransfer` model.

    See Also
    --------
    StyleTransfer

    Examples
    --------
    .. sourcecode:: python

        # Create datasets
        >>> content_dataset = turicreate.image_analysis.load_images('content_images/')
        >>> style_dataset = turicreate.image_analysis.load_images('style_images/')

        # Train a style transfer model
        >>> model = turicreate.style_transfer.create(content_dataset, style_dataset)

        # Stylize an image on all styles
        >>> stylized_images = model.stylize(data)

        # Visualize the stylized images
        >>> stylized_images.explore()

    """
    if len(style_dataset) == 0:
        raise _ToolkitError("style_dataset SFrame cannot be empty")
    if len(content_dataset) == 0:
        raise _ToolkitError("content_dataset SFrame cannot be empty")
    if(batch_size < 1):
        raise _ToolkitError("'batch_size' must be greater than or equal to 1")

    from ._sframe_loader import SFrameSTIter as _SFrameSTIter
    import mxnet as _mx
    from .._mxnet import _mxnet_utils

    if style_feature is None:
        style_feature = _tkutl._find_only_image_column(style_dataset)
    if content_feature is None:
        content_feature = _tkutl._find_only_image_column(content_dataset)
    if verbose:
        print("Using '{}' in style_dataset as feature column and using "
              "'{}' in content_dataset as feature column".format(style_feature, content_feature))

    _raise_error_if_not_training_sframe(style_dataset, style_feature)
    _raise_error_if_not_training_sframe(content_dataset, content_feature)

    params = {
        'batch_size': batch_size,
        'vgg16_content_loss_layer': 2,  # conv3_3 layer
        'lr': 0.001,
        'content_loss_mult': 1.0,
        'style_loss_mult': [1e-4, 1e-4, 1e-4, 1e-4],  # conv 1-4 layers
        'finetune_all_params': True,
        'pretrained_weights': False,
        'print_loss_breakdown': False,
        'input_shape': (256, 256),
        'training_content_loader_type': 'stretch',
        'use_augmentation': False,
        'sequential_image_processing': False,
        # Only used if use_augmentaion is True
        'aug_resize': 0,
        'aug_min_object_covered': 0,
        'aug_rand_crop': 0.9,
        'aug_rand_pad': 0.9,
        'aug_rand_gray': 0.0,
        'aug_aspect_ratio': 1.25,
        'aug_hue': 0.05,
        'aug_brightness': 0.05,
        'aug_saturation': 0.05,
        'aug_contrast': 0.05,
        'aug_horizontal_flip': True,
        'aug_area_range': (.05, 1.5),
        'aug_pca_noise': 0.0,
        'aug_max_attempts': 20,
        'aug_inter_method': 2,
    }

    if '_advanced_parameters' in kwargs:
        # Make sure no additional parameters are provided
        new_keys = set(kwargs['_advanced_parameters'].keys())
        set_keys = set(params.keys())
        unsupported = new_keys - set_keys
        if unsupported:
            raise _ToolkitError('Unknown advanced parameters: {}'.format(unsupported))

        params.update(kwargs['_advanced_parameters'])

    _content_loss_mult = params['content_loss_mult']
    _style_loss_mult = params['style_loss_mult']

    num_gpus = _mxnet_utils.get_num_gpus_in_use(max_devices=params['batch_size'])
    batch_size_each = params['batch_size'] // max(num_gpus, 1)
    batch_size = max(num_gpus, 1) * batch_size_each
    input_shape = params['input_shape']

    iterations = 0
    if max_iterations is None:
        max_iterations = len(style_dataset) * 10000
        if verbose:
            print('Setting max_iterations to be {}'.format(max_iterations))

    # data loader
    if params['use_augmentation']:
        content_loader_type = '%s-with-augmentation' % params['training_content_loader_type']
    else:
        content_loader_type = params['training_content_loader_type']

    content_images_loader = _SFrameSTIter(content_dataset, batch_size, shuffle=True,
                                  feature_column=content_feature, input_shape=input_shape,
                                  loader_type=content_loader_type, aug_params=params,
                                  sequential=params['sequential_image_processing'])
    ctx = _mxnet_utils.get_mxnet_context(max_devices=params['batch_size'])

    num_styles = len(style_dataset)

    # TRANSFORMER MODEL
    from ._model import Transformer as _Transformer
    transformer_model_path = _pre_trained_models.STYLE_TRANSFER_BASE_MODELS[model]().get_model_path()
    transformer = _Transformer(num_styles, batch_size_each)
    transformer.collect_params().initialize(ctx=ctx)
    
    if params['pretrained_weights']:
        transformer.load_params(transformer_model_path, ctx, allow_missing=True)

    # For some reason, the transformer fails to hybridize for training, so we
    # avoid this until resolved
    # transformer.hybridize()

    # VGG MODEL
    from ._model import Vgg16 as _Vgg16
    vgg_model_path = _pre_trained_models.STYLE_TRANSFER_BASE_MODELS['Vgg16']().get_model_path()
    vgg_model = _Vgg16()
    vgg_model.collect_params().initialize(ctx=ctx)
    vgg_model.load_params(vgg_model_path, ctx=ctx, ignore_extra=True)
    vgg_model.hybridize()

    # TRAINER
    from mxnet import gluon as _gluon
    from ._model import gram_matrix as _gram_matrix

    if params['finetune_all_params']:
        trainable_params = transformer.collect_params()
    else:
        trainable_params = transformer.collect_params('.*gamma|.*beta')

    trainer = _gluon.Trainer(trainable_params, 'adam', {'learning_rate': params['lr']})
    mse_loss = _gluon.loss.L2Loss()
    start_time = _time.time()
    smoothed_loss = None
    last_time = 0

    cuda_gpus = _mxnet_utils.get_gpus_in_use(max_devices=params['batch_size'])
    num_mxnet_gpus = len(cuda_gpus)

    if verbose:
        # Estimate memory usage (based on experiments)
        cuda_mem_req = 260 + batch_size_each * 880 + num_styles * 1.4

        _tkutl._print_neural_compute_device(cuda_gpus=cuda_gpus, use_mps=False,
                                            cuda_mem_req=cuda_mem_req, has_mps_impl=False)
    #
    # Pre-compute gram matrices for style images
    #
    if verbose:
        print('Analyzing visual features of the style images')

    style_images_loader = _SFrameSTIter(style_dataset, batch_size, shuffle=False, num_epochs=1,
                                        feature_column=style_feature, input_shape=input_shape,
                                        loader_type='stretch',
                                        sequential=params['sequential_image_processing'])
    num_layers = len(params['style_loss_mult'])
    gram_chunks = [[] for _ in range(num_layers)]
    for s_batch in style_images_loader:
        s_data = _gluon.utils.split_and_load(s_batch.data[0], ctx_list=ctx, batch_axis=0)
        for s in s_data:
            vgg16_s = _vgg16_data_prep(s)
            ret = vgg_model(vgg16_s)
            grams = [_gram_matrix(x) for x in ret]
            for i, gram in enumerate(grams):
                if gram.context != _mx.cpu(0):
                    gram = gram.as_in_context(_mx.cpu(0))
                gram_chunks[i].append(gram)
    del style_images_loader

    grams = [
        # The concatenated styles may be padded, so we slice overflow
        _mx.nd.concat(*chunks, dim=0)[:num_styles]
        for chunks in gram_chunks
    ]

    # A context->grams look-up table, where all the gram matrices have been
    # distributed
    ctx_grams = {}
    if ctx[0] == _mx.cpu(0):
        ctx_grams[_mx.cpu(0)] = grams
    else:
        for ctx0 in ctx:
            ctx_grams[ctx0] = [gram.as_in_context(ctx0) for gram in grams]

    #
    # Training loop
    #

    vgg_content_loss_layer = params['vgg16_content_loss_layer']
    rs = _np.random.RandomState(1234)
    while iterations < max_iterations:
        content_images_loader.reset()
        for c_batch in content_images_loader:
            c_data = _gluon.utils.split_and_load(c_batch.data[0], ctx_list=ctx, batch_axis=0)

            Ls = []
            curr_content_loss = []
            curr_style_loss = []
            with _mx.autograd.record():
                for c in c_data:
                    # Randomize styles to train
                    indices = _mx.nd.array(rs.randint(num_styles, size=batch_size_each),
                                           dtype=_np.int64, ctx=c.context)

                    # Generate pastiche
                    p = transformer(c, indices)

                    # mean subtraction
                    vgg16_p = _vgg16_data_prep(p)
                    vgg16_c = _vgg16_data_prep(c)

                    # vgg forward
                    p_vgg_outputs = vgg_model(vgg16_p)

                    c_vgg_outputs = vgg_model(vgg16_c)
                    c_content_layer = c_vgg_outputs[vgg_content_loss_layer]
                    p_content_layer = p_vgg_outputs[vgg_content_loss_layer]

                    # Calculate Loss
                    # Style Loss between style image and stylized image
                    # Ls = sum of L2 norm of gram matrix of vgg16's conv layers
                    style_losses = []
                    for gram, p_vgg_output, style_loss_mult in zip(ctx_grams[c.context], p_vgg_outputs, _style_loss_mult):
                        gram_s_vgg = gram[indices]
                        gram_p_vgg = _gram_matrix(p_vgg_output)

                        style_losses.append(style_loss_mult * mse_loss(gram_s_vgg, gram_p_vgg))

                    style_loss = _mx.nd.add_n(*style_losses)

                    # Content Loss between content image and stylized image
                    # Lc = L2 norm at a single layer in vgg16
                    content_loss = _content_loss_mult * mse_loss(c_content_layer,
                                                                 p_content_layer)

                    curr_content_loss.append(content_loss)
                    curr_style_loss.append(style_loss)
                    # Divide loss by large number to get into a more legible
                    # range
                    total_loss = (content_loss + style_loss) / 10000.0
                    Ls.append(total_loss)
                for L in Ls:
                    L.backward()

            cur_loss = _np.mean([L.asnumpy()[0] for L in Ls])

            if smoothed_loss is None:
                smoothed_loss = cur_loss
            else:
                smoothed_loss = 0.9 * smoothed_loss + 0.1 * cur_loss
            iterations += 1
            trainer.step(batch_size)

            if verbose and iterations == 1:
                # Print progress table header
                column_names = ['Iteration', 'Loss', 'Elapsed Time']
                num_columns = len(column_names)
                column_width = max(map(lambda x: len(x), column_names)) + 2
                hr = '+' + '+'.join(['-' * column_width] * num_columns) + '+'
                print(hr)
                print(('| {:<{width}}' * num_columns + '|').format(*column_names, width=column_width-1))
                print(hr)

            cur_time = _time.time()
            if verbose and (cur_time > last_time + 10 or iterations == max_iterations):
                # Print progress table row
                elapsed_time = cur_time - start_time
                print("| {cur_iter:<{width}}| {loss:<{width}.3f}| {time:<{width}.1f}|".format(
                    cur_iter = iterations, loss = smoothed_loss,
                    time = elapsed_time , width = column_width-1))
                if params['print_loss_breakdown']:
                    print_content_loss = _np.mean([L.asnumpy()[0] for L in curr_content_loss])
                    print_style_loss = _np.mean([L.asnumpy()[0] for L in curr_style_loss])
                    print('Total Loss: {:6.3f} | Content Loss: {:6.3f} | Style Loss: {:6.3f}'.format(cur_loss, print_content_loss, print_style_loss))
                last_time = cur_time
            if iterations == max_iterations:
                print(hr)
                break

    training_time = _time.time() - start_time
    style_sa = style_dataset[style_feature]
    idx_column = _tc.SArray(range(0, style_sa.shape[0]))
    style_sframe = _tc.SFrame({"style": idx_column, style_feature: style_sa})

    # Save the model state
    state = {
        '_model': transformer,
        '_training_time_as_string': _seconds_as_string(training_time),
        'batch_size': batch_size,
        'num_styles': num_styles,
        'model': model,
        'input_image_shape': input_shape,
        'styles': style_sframe,
        'num_content_images': len(content_dataset),
        'training_time': training_time,
        'max_iterations': max_iterations,
        'training_iterations': iterations,
        'training_epochs': content_images_loader.cur_epoch,
        'style_feature': style_feature,
        'content_feature': content_feature,
        "_index_column": "style",
        'training_loss': smoothed_loss,
    }

    return StyleTransfer(state)