app.py

import numpy as np
import tensorflow.lite as tfl
import matplotlib.pyplot as plt
import torchvision.utils as vutils
import torch
import gradio as gr

print("imported stuff")
interpreter = tfl.Interpreter(model_path="trainemall.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] 
  
  plt.figure(figsize=(7,7))
  plt.imshow(np.transpose(first_image.numpy(), (1, 2, 0)))
  mon = plt.axis('off')  
  plt.gcf()

demo = gr.Interface(
    makemon,
    [
    ],
    outputs = gr.Plot(label="Pokemon", format="png"),
)

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