import gradio as gr
import time
import keras_cv
from tensorflow import keras
import matplotlib.pyplot as plt
from translate import Translator

keras.mixed_precision.set_global_policy("mixed_float16")
model = keras_cv.models.StableDiffusion(img_width=256, img_height=256, jit_compile=True)

def plot_images(images):
    plt.figure(figsize=(10, 10))
    for i in range(len(images)):
        ax = plt.subplot(1, len(images), i + 1)
        plt.imshow(images[i])
        plt.axis("off")
        plt.tight_layout()

translator = Translator(from_lang="ko", to_lang="en")

def generate_images(text, n=3):
    print(text)
    translation = translator.translate(text)
    print(translation)
    images = model.text_to_image(translation, batch_size=n)
    return images

def inference(text):
    image = generate_images(text, 1).squeeze()
    return image

demo = gr.Interface(fn=inference, inputs="text", outputs="image")

demo.launch(share=True)