import gradio as gr

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe = pipe.to("mps")

# Recommended if your computer has < 64 GB of RAM
pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
#image = pipe(prompt).images[0]


text0 = 'black fluffy gorgeous dangerous cat animal creature, large orange eyes, big fluffy ears, piercing gaze, full moon, dark ambiance, best quality, extremely detailed'

def generation(text):
    image = pipe(text).images[0]
    return image

demo = gr.Blocks()

title = '# 3D print failures detection App'
description = 'App for detect errors in the 3D printing'

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Row():
        img_input = gr.Textbox ( label="Text 1",info="Initial text",lines=5,value=text0)

        button = gr.Button(value="Generate")

        
    with gr.Row():
        img_output= gr.Image()
    
    button.click( generation, inputs=img_input, outputs=[img_output])
        


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