app.py

import streamlit as st
import torch
import io
from diffusers import StableDiffusionPipeline
from PIL import Image

# Disable unnecessary warnings
import warnings
warnings.filterwarnings("ignore")

# Model configuration
MODEL_ID = "CompVis/stable-diffusion-v1-4"

@st.cache_resource(show_spinner=False)
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        use_safetensors=False,  # Force using .bin files
        use_auth_token=False,
        safety_checker=None
    )
    pipe = pipe.to("cpu")
    pipe.enable_attention_slicing()
    return pipe

def main():
    st.set_page_config(page_title="Poetry to Image", page_icon="🎨")
    st.title("✨ Romantic Text to Image Converter")
    
    # Initialize session state
    if 'generated_image' not in st.session_state:
        st.session_state.generated_image = None
    if 'error' not in st.session_state:
        st.session_state.error = None

    # Sidebar controls
    with st.sidebar:
        st.header("Settings")
        num_steps = st.slider("Inference Steps", 10, 50, 25)
        guidance_scale = st.slider("Guidance Scale", 3.0, 15.0, 7.5)

    # Main interface
    prompt = st.text_area(
        "Enter your romantic text:",
        height=150,
        placeholder="Example: 'Your smile shines brighter than the sun...'"
    )

    if st.button("Generate Image", type="primary"):
        st.session_state.generated_image = None
        st.session_state.error = None
        
        if not prompt.strip():
            st.session_state.error = "Please enter some text first!"
        else:
            try:
                with st.spinner("Creating artwork (5-10 minutes)..."):
                    pipe = load_model()
                    image = pipe(
                        prompt,
                        num_inference_steps=num_steps,
                        guidance_scale=guidance_scale
                    ).images[0]
                    st.session_state.generated_image = image
            except Exception as e:
                st.session_state.error = f"Error: {str(e)}"

    # Display results
    if st.session_state.error:
        st.error(st.session_state.error)
    
    if st.session_state.generated_image:
        st.subheader("Generated Image")
        st.image(st.session_state.generated_image, use_column_width=True)
        
        # Download functionality
        buf = io.BytesIO()
        st.session_state.generated_image.save(buf, format="PNG")
        st.download_button(
            "Download Image",
            buf.getvalue(),
            "poetry_image.png",
            "image/png"
        )

if __name__ == "__main__":
    main()