import gradio as gr
from PIL import Image  
from diffusers import StableDiffusionPipeline


katz_model_path = 'maximalmargin/katz'
katz = StableDiffusionPipeline.from_pretrained(katz_model_path)


def predict(text):
    image = katz(prompt = text).images[0]
    return [image]

gallery = gr.Gallery(
    label="Generated images", show_label=False, elem_id="gallery"
).style(height="auto")

gr.Interface(
    predict,
    inputs=gr.Textbox(
                        label="Enter your prompt",
                        show_label=False,
                        max_lines=1,
                        placeholder="Enter your prompt",
                    ),
    outputs=[gallery],
    title="Alex Katz Portraiter",
).launch()