2DFloors / app.py
sialnoman318's picture
Update app.py
fbd403f verified
import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
import transformers
# Load Stable Diffusion model
model_id = "stabilityai/stable-diffusion-v1-4"
model = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
# Streamlit UI
st.title("Generative AI Application")
# Input prompt
def main():
st.sidebar.header("Input Settings")
prompt = st.sidebar.text_input("Enter your prompt:", "A modern house floor plan with 3 bedrooms and a pool")
if st.sidebar.button("Generate Image"):
if not prompt:
st.error("Please enter a prompt.")
else:
with st.spinner("Generating image..."):
try:
images = model([prompt])
image = images[0]
st.image(image, caption="Generated Image", use_column_width=True)
except Exception as e:
st.error(f"Error generating image: {str(e)}")
if __name__ == "__main__":
main()