app.py

import torch
import streamlit as st
from diffusers import StableDiffusion3Pipeline

# Retry mechanism for loading the model
def retry_load_model():
    try:
        return StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
    except EnvironmentError as e:
        st.error(f"Error loading model: {e}")
        st.warning("Retrying to load the model...")
        retry_load_model()

# Load the Diffusion pipeline
pipeline = retry_load_model()

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()