Ccre/SD3_WebUi

1import streamlit as st2import requests3from PIL import Image4import datetime5import io6import os7 8output_dir = "outputs"9os.makedirs(output_dir, exist_ok=True)10 11st.set_page_config(initial_sidebar_state="collapsed", page_icon="static/sai_logo.ico", layout="centered")12 13def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_format, strength=None, negative_prompt=None, image_file=None): #, num_outputs=1):14    headers = {15        "authorization": f"Bearer {api_key}",16        "accept": "image/*"17    }18    data = {19        "prompt": prompt,20        "mode": mode,21        "model": model,22#       "num_outputs": num_outputs,23        "output_format": output_format24    }25    26    files = {"none": ''}  # Default to no file if none provided27 28    if mode == "text-to-image":29        files["aspect_ratio"] = (None, aspect_ratio)30    elif mode == "image-to-image":31        if seed != 0:32            data['seed'] = seed33    if negative_prompt:34            data['negative_prompt'] = negative_prompt35    if image_file is not None:36            image_bytes = image_file.getvalue()37            files = {'image': ('image.jpg', image_bytes, 'image/jpeg')}38    if strength is not None:39        files["strength"] = (None, str(strength))40    if seed is not None:41        files["seed"] = (None, str(seed))42    response = requests.post(43        "https://api.stability.ai/v2beta/stable-image/generate/sd3",44        headers=headers,45        files=files,46        data=data47    )48    return response49 50def main():51    st.image("static/SD3_webui_logo_image.png", use_column_width=True)52    53    if 'api_key' not in st.session_state:54        st.session_state.api_key = ""55    if 'prompt' not in st.session_state:56        st.session_state.prompt = ""57    if 'negative_prompt' not in st.session_state:58        st.session_state.negative_prompt = ""59    if 'seed' not in st.session_state:60        st.session_state.seed = 061    62    api_key = st.text_input("Enter your API key:", type="password", value=st.session_state.api_key)63    models = st.selectbox("Select model:", ["sd3", "sd3-turbo"], index=0)64    mode = st.selectbox("Select generation mode:", ["text-to-image", "image-to-image"], index=0)65    prompt = st.text_area("Enter positive prompt:", value=st.session_state.prompt)66    negative_prompt = st.text_input("Enter negative prompt: (optional)", value=st.session_state.negative_prompt)67    aspect_ratios = st.selectbox("Select the aspect ratio:", ["1:1", "2:3", "3:2", "4:5", "5:4", "16:9", "21:9", "9:16", "9:21"], index=0) if mode == "text-to-image" else None68    seed = st.slider("Set seed: (randomize = 0)", min_value=0, max_value=4294967294, value=st.session_state.seed, step=1)69    strength = st.slider("Denoising strength:", min_value=0.0, max_value=1.0, value=0.5, step=0.01) if mode == "image-to-image" else None70#   num_outputs = st.number_input("Batch Count:", min_value=1, max_value=20, value=1, step=1)71    output_formats = st.selectbox("Select the output format:", ["jpeg", "png"], index=1)72    73    st.session_state.api_key = api_key74    st.session_state.prompt = prompt75    st.session_state.negative_prompt = negative_prompt76    st.session_state.seed = seed77    78    uploaded_file = st.file_uploader("Upload for image-to-image generation:", type=["png", "jpg", "jpeg"]) if mode == "image-to-image" else None79 80    if st.button('Generate Image', use_container_width=True):81        with st.spinner('Generating Image...'):82            result = generate_image(api_key, prompt, aspect_ratios, mode, models, seed, output_formats, strength, negative_prompt, uploaded_file)83            if result.status_code == 200:84                image = Image.open(io.BytesIO(result.content))85                st.image(image, use_column_width=True)86                current_time = datetime.datetime.now().strftime("%Y%m%d_%H.%M.%S").lower()87                file_path = os.path.join(output_dir, f"gen_{models}_{seed}_{current_time}.{output_formats}")88                image.save(file_path)89                st.success(f'Image saved to {file_path}')90            else:91                st.error('Failed to generate image: ' + str(result.json()))92                93    with st.sidebar:94        st.link_button("Stability.Ai _ API Documentation", "https://platform.stability.ai/docs/api-reference")95        with st.expander(("Image Generation Costs"), expanded=True):96            st.markdown((97            """98            - SD3 - 6.5 credits per image or $0.06599            - SD3 Turbo - 4.0 credits per image or $0.04100            101            Additional credits can be purchased via the account_page button below.102            103            Credits cost $10 per 1,000 credits, which is    enough credits for roughly 154 SD3 images or 250 SD3 Turbo images.104            """105        ))106        st.link_button("Account_Page", "https://platform.stability.ai/account/credits")107                108if __name__ == "__main__":109    main()