snehakingrani's picture
Update app.py
2381e5b verified
raw
history blame contribute delete
729 Bytes
import streamlit as st
from PIL import Image
import torch
from torchvision import transforms
from bk_sdm import BKSDMPipeline # Hypothetical import; replace with actual model import
# Initialize the model
@st.cache_resource
def load_model():
model = BKSDMPipeline.from_pretrained('path_to_bk_sdm_model')
return model
model = load_model()
# Streamlit app interface
st.title("Text-to-Image Generation")
prompt = st.text_input("Enter a text prompt:")
if st.button("Generate Image"):
if prompt:
with st.spinner("Generating image..."):
# Generate image
image = model(prompt)
st.image(image, caption="Generated Image")
else:
st.warning("Please enter a prompt.")