from cProfile import label
import glob
import gradio as gr
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
import numpy as np

pixel_cnn = from_pretrained_keras("keras-io/pixel-cnn-mnist")


def infer(batch):
    pixels = np.zeros(shape=(batch,) + (pixel_cnn.input_shape)[1:])
    batch, rows, cols, channels = pixels.shape

    # Iterate over the pixels because generation has to be done sequentially pixel by pixel.
    for row in range(rows):
        for col in range(cols):
            for channel in range(channels):
                # Feed the whole array and retrieving the pixel value probabilities for the next
                # pixel.
                probs = pixel_cnn.predict(pixels)[:, row, col, channel]
                # Use the probabilities to pick pixel values and append the values to the image
                # frame.
                pixels[:, row, col, channel] = tf.math.ceil(
                    probs - tf.random.uniform(probs.shape)
                )

    for i, pic in enumerate(pixels):
        tf.keras.preprocessing.image.save_img(
            "/tmp/generated_image_{}.png".format(i), deprocess_image(np.squeeze(pic, -1))
        )
    return glob.glob("/tmp/generated*")


def deprocess_image(x):
    # Stack the single channeled black and white image to RGB values.
    x = np.stack((x, x, x), 2)
    # Undo preprocessing
    x *= 255.0
    # Convert to uint8 and clip to the valid range [0, 255]
    x = np.clip(x, 0, 255).astype("uint8")
    return x


article = """<center>
Authors: Space by <a href='https://twitter.com/johko990' target='_blank'><b>Johannes Kolbe</b></a>, model by İhsan Soydemir after an example by ADMoreau at 
<a href='https://keras.io/examples/generative/pixelcnn/' target='_blank'><b>keras.io</b></a> <br>
<a href='https://arxiv.org/abs/1606.05328' target='_blank'><b>Original paper</b></a> by van den Oord et al."""


description = """Image generation using a CNN. The model is trained on MNIST data, so the generation capabilities are limited to MNIST like images. <br>
Just use the slider to set how many images should be created.<br>
The execution might take some time."""


Iface = gr.Interface(
    fn=infer,
    inputs=gr.inputs.Slider(minimum=1, maximum=20, default=4, step=1, label="Number of images to generate"),
    outputs=gr.outputs.Carousel(["image"]),
    title="PixelCNN - MNIST Image Generation",
    article=article,
    description=description,
).launch(enable_queue=True)