# app.py

import streamlit as st
import os
from getpass import getpass
from replicate import replicate
from IPython.display import Image

# Set up Replicate API token
st.title('Stable Diffusion 3 Image Generation')

if 'REPLICATE_API_TOKEN' not in os.environ:
    st.warning('Please set your Replicate API token. You can obtain it from https://replicate.com/account.')
    REPLICATE_API_TOKEN = st.text_input('Enter your Replicate API token:')
    if REPLICATE_API_TOKEN:
        os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN
else:
    st.write('API token is set.')

# Function to generate image
@st.cache(allow_output_mutation=True)
def generate_image():
    output = replicate.run(
        "stability-ai/stable-diffusion-3",
        input={
            "cfg": 4.5,
            "prompt": "a photo of vibrant artistic graffiti on a wall saying \"SD3 medium\"",
            "aspect_ratio": "3:2",
            "output_format": "webp",
            "output_quality": 79,
            "negative_prompt": "ugly, distorted"
        }
    )
    return output[0]

# Main Streamlit app
def main():
    st.subheader('Generate Image')
    if 'REPLICATE_API_TOKEN' in os.environ:
        image_url = generate_image()
        st.image(image_url)
    else:
        st.warning('Please set your Replicate API token to generate an image.')

if __name__ == '__main__':
    main()