retrAIced / app.py
JavierGon12's picture
Add custom prompt
86165e8
raw
history blame
403 Bytes
import streamlit as st
from diffusers import DDPMScheduler, UNet2DModel
from PIL import Image
import torch
import numpy as np
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
prompt = st.text_input('Insert here your prompt')
image = pipe(prompt).images[0]