File size: 617 Bytes
7c373db
 
54d6e13
4430ce1
 
a543549
375ae93
 
4430ce1
7c373db
 
 
 
4430ce1
 
 
7c373db
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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")