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>
    """
inputs = gr.inputs.Number(label="number of images")
outputs = gr.outputs.Image(label="Predictions")

def create_digit_samples(n_samples):
      latent_dim = 128
      random_latent_vectors = tf.random.normal(shape=(int(n_samples), 128))
      examples = model.predict(random_latent_vectors)
      #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


inputs = gr.inputs.Number(label="number of images")
outputs = gr.outputs.Image(label="Output Image")

examples = [
            [1], 
            [2],
            [3],
            [4], 
            [64]
]


gr.Interface(create_digit_samples, inputs, outputs, analytics_enabled=False, examples=examples).launch(debug=True)