import torch
import streamlit as st
from diffusers import StableDiffusion3Pipeline

# Load the Diffusion pipeline
pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)

# Your other functions and main code here...


def generate_prompt(prompt_text):
    # Generate response using the Diffusion model
    response = pipeline(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()