import streamlit as st
from diffusers import DiffusionPipeline

# Load the Diffusion pipeline
@st.cache(allow_output_mutation=True)
def load_diffusion_pipeline():
    try:
        pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium")
        return pipeline
    except Exception as e:
        st.error(f"Error loading model: {e}")

def generate_response(prompt_text, pipeline):
    try:
        response = pipeline(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
        return response
    except Exception as e:
        st.error(f"Error generating response: {e}")

def main():
    st.title('Hugging Face Diffusion Model')

    # Load the model
    pipeline = load_diffusion_pipeline()

    # Text input for the prompt
    prompt_text = st.text_area("Enter your prompt here:", height=200)

    # Button to generate prompt
    if st.button("Generate"):
        if prompt_text:
            with st.spinner('Generating...'):
                generated_text = generate_response(prompt_text, pipeline)
            st.success('Generation complete!')
            st.text_area('Generated Text:', value=generated_text, height=400)
        else:
            st.warning('Please enter a prompt.')

if __name__ == '__main__':
    main()