import numpy as np
import tensorflow.lite as tfl
import matplotlib.pyplot as plt
import gradio as gr
import torchvision.utils as vutils
import torch
print("imported stuff")
interpreter = tfl.Interpreter(model_path="resnet.tflite")
interpreter.allocate_tensors()
print('started up model')
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

def makemon():
  input_data = np.array(np.random.normal(size=input_shape), dtype=np.float32)
  interpreter.set_tensor(input_details[0]['index'], input_data)
  
  interpreter.invoke()
  
  output_data = interpreter.get_tensor(output_details[0]['index'])
  
  grid = vutils.make_grid(torch.from_numpy(output_data).detach().cpu(), padding=0, normalize=True)
  
  first_image = grid[:, :grid.size(1) // 8, :grid.size(2) // 8]  # Assuming an 8x8 grid
  
  plt.figure(figsize=(7, 7))
  plt.imshow(np.transpose(first_image.numpy(), (1, 2, 0)))
  mon = plt.axis('off')  # Turn off axis for a cleaner look
  return plt.gcf()
a = makemon()
print("mademon")
demo = gr.Interface(
    makemon,
    [
    ],
    outputs = gr.Plot(label="Pokemon", format="png"),
)

if __name__ == "__main__":
    demo.launch(debug=True)