rainbow-canvas / app.py
homer-meng's picture
Update app.py
375ae93
raw
history blame
617 Bytes
import streamlit as st
import torch
from transformers import AutoTokenizer, GPTNeoForCausalLM
from PIL import Image
import io
tokenizer = AutoTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5")
model = GPTNeoForCausalLM.from_pretrained("runwayml/stable-diffusion-v1-5")
st.title("Scribble Drawing Generator")
prompt = st.text_input("Enter a prompt:")
if prompt:
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(inputs['input_ids'], max_length=256, do_sample=True)
image = Image.open(io.BytesIO(outputs[0].cpu().numpy()))
st.image(image, caption="Generated Drawing")