import streamlit as st
from transformers import pipeline

# Load the Diffusion pipeline for text generation
generator = pipeline("text-generation", model="stabilityai/stable-diffusion-3-medium")

def generate_prompt(prompt_text):
    # Generate response using the Diffusion model
    response = generator(prompt_text, top_p=0.9, max_length=100)[0]['generated_text']
    return response

def main():
    st.title('Diffusion Model Prompt Generator')

    # 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_prompt(prompt_text)
            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()