from stable_diffusion_tf.stable_diffusion import Text2Image
from PIL import Image
import gradio as gr

generator = Text2Image(
    img_height=512,
    img_width=512,
    jit_compile=False,
)

def txt2img(prompt, guide, steps, Temp):
    img = generator.generate(prompt,
    num_steps=steps,
    unconditional_guidance_scale=guide,
    temperature=Temp,
    batch_size=1)
    image=Image.fromarray(img[0])
    return image

iface = gr.Interface(fn=txt2img, inputs=[
    gr.Textbox(label = 'Input Text Prompt'),
    gr.Slider(2, 20, value = 9, label = 'Guidence Scale'),
    gr.Slider(10, 100, value = 50, step = 1, label = 'Number of Iterations'),
    gr.Slider(1, 100, value=1, step=1)], outputs = 'image',title='Stable Diffusion with Keras and TensorFlow CPU or GPU', description='Now Using Keras and TensorFlow with Stable Diffusion. This allows very complex image generation with less code footprint, and less text.')
iface.launch()