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")