app.py

import gradio as gr
import torch

from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"
# if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)

def generate(prompt):
  prefix = 'A painting of'
  prefix += prompt
  add = '#artstation'
  prefix += add
  true_prompt = prefix

  prompt = "a photo of an astronaut riding a horse on mars"
  image = pipe(true_prompt).images[0]  
    
  image.save("result.png")
  return image

title = 'Art generator'

description = 'Input a prompt and generate art inspired from it!'

ui = gr.Interface(generate, inputs=['text'], outputs=['image'], description=description, title=title)
ui.launch()