from huggingface_hub import from_pretrained_keras
import matplotlib.pyplot as plt
from math import sqrt, ceil
import tensorflow as tf
import gradio as gr
import numpy as np

model = from_pretrained_keras("IMvision12/WGAN-GP")

title = "WGAN-GP"
description = "Image Generation Using WGAN"
article = """
<p style='text-align: center'>
            <a href='https://keras.io/examples/generative/wgan_gp/' target='_blank'>Keras Example given by A_K_Nain</a>
            <br>
            Space by Gitesh Chawda
        </p>
    """

def generate_latent_points(latent_dim, n_samples):
	random_latent_vectors = tf.random.normal(shape=(n_samples, latent_dim))
	return random_latent_vectors

def create_digit_samples(n_samples):
        latent_dim = 128
        random_vector_labels = generate_latent_points(latent_dim, int(n_samples))
        examples = model.predict(random_vector_labels)
        examples = examples * 255.0
        size = ceil(sqrt(n_samples))
        digit_images = np.zeros((28*size, 28*size), dtype=float)
        n = 0
        for i in range(size):
            for j in range(size):
                if n == n_samples:
                    break
                digit_images[i* 28 : (i+1)*28, j*28 : (j+1)*28] = examples[n, :, :, 0]
                n += 1
        digit_images = (digit_images/127.5) -1
        return digit_images


examples = [[1, 10], [3, 5], [5, 15]]


iface = gr.Interface(
    fn = create_digit_samples,
    inputs = ["number", "number"],
    outputs = ["image"],
    examples = examples,
    description = description,
    title = title,
    article = article
    )

iface.launch()