# notes https://huggingface.co/spaces/Joeythemonster/Text-To-image-AllModels/blob/main/app.py

import tensorflow as tf
from diffusers import DiffusionPipeline
import spaces
# import torch
import PIL.Image
import gradio as gr
import gradio.components as grc
import numpy as np
from huggingface_hub import from_pretrained_keras
import keras
import time
import os 

# os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

# options = ['Placeholder A', 'Placeholder B', 'Placeholder C']


# pipeline = DiffusionPipeline.from_pretrained("nathanReitinger/MNIST-diffusion-oneImage")
# device = "cuda" if torch.cuda.is_available() else "cpu"
# pipeline = pipeline.to(device=device)

# @spaces.GPU
# def predict(steps, seed):
#		 print("HI")
#		 generator = torch.manual_seed(seed)
#		 for i in range(1,steps):
#				 yield pipeline(generator=generator, num_inference_steps=i).images[0]

# gr.Interface(
#		 predict,
#		 inputs=[
#				 grc.Slider(0, 1000, label='Inference Steps', value=42, step=1),
#				 grc.Slider(0, 2147483647, label='Seed', value=42, step=1),
#		 ],
#		 outputs=gr.Image(height=28, width=28, type="pil", elem_id="output_image"),
#		 css="#output_image{width: 256px !important; height: 256px !important;}",
#		 title="Model Problems: Infringing on MNIST!",
#		 description="Opening the black box.",
# ).queue().launch()


from diffusers import StableDiffusionPipeline
import torch


modellist=['nathanReitinger/MNIST-diffusion-oneImage',
	'nathanReitinger/MNIST-diffusion',
#	'nathanReitinger/MNIST-GAN', 
#	'nathanReitinger/MNIST-GAN-noDropout'
 ]

# pipeline = DiffusionPipeline.from_pretrained("nathanReitinger/MNIST-diffusion-oneImage")
# device = "cuda" if torch.cuda.is_available() else "cpu"
# pipeline = pipeline.to(device=device)


def getModel(model):
	model_id = model

	(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()
	RANDO = str(time.time())
	file_path = 'tester/' + model_id.replace("/", "-") + "/" + RANDO + '/'
	os.makedirs(file_path)
	train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
	train_images = (train_images - 127.5) / 127.5	# Normalize the images to [-1, 1]

	print(model_id)
	if 'diffusion' in model_id:
		pipe = DiffusionPipeline.from_pretrained(model_id)
		pipe = pipe.to("cpu")
		image = pipe(generator= torch.manual_seed(42), num_inference_steps=40).images[0]
	else:
		pipe = DiffusionPipeline.from_pretrained('nathanReitinger/MNIST-diffusion')
		pipe = pipe.to("cpu")
		test = from_pretrained_keras('nathanReitinger/MNIST-GAN')
		image = pipe(generator= torch.manual_seed(42), num_inference_steps=40).images[0]

	return image


import gradio as gr
interface = gr.Interface(fn=getModel, inputs=[gr.Dropdown(modellist)], css="#output_image{width: 256px !important; height: 256px !important;}", outputs="image",  title='Model Problems (infringement)')
interface.launch()