import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.python.ops.numpy_ops import np_config
from visualblocks import register_vb_fn, Server

np_config.enable_numpy_behavior()

hub_handle = "https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2"
hub_module = hub.load(hub_handle)


# Register the function with visual blocks using the "generic" type (meaning
# tensors in, tensors out)
@register_vb_fn(type="generic")
def styleTransfer(tensors):
    """Inference function for use with Visual Blocks.

    This function is passed to the Visual Blocks server, which calls it to
    implement a Colab model runner block.

    Args:
      tensors: A list of np.ndarrays as input tensors. For this particular
      inference function, only the first two np.ndarrays are used. The first
      np.ndarrays is the input content image as a tensor of size [1,
      content_image_height, content_image_width, 3] with floating point pixel
      values ranging from 0 to 1. The second np.ndarrays is the
      input style image as a tensor of size [1, style_image_height,
      style_image_width, 3] with floating point pixel values ranging from 0 to 1.

    Returns:
      tensors: A list of np.ndarrays as output tensors. For this particular
      inference function, only the first item is used. The first item is the
      output image as a tensor of size [1, height, width, 3] with floating point
      pixel values ranging from 0 to 1.
    """

    content_tensor = tf.constant(tensors[0], dtype=tf.float32)
    style_tensor = tf.constant(tensors[1], dtype=tf.float32)
    outputs = hub_module(content_tensor, style_tensor)
    stylized_image = outputs[0].numpy()

    return [
        stylized_image,
    ]


server = Server()
server.run()