app.py

import gradio as gr
from stability_sdk import client
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
from PIL import Image
import io
import os
import warnings
from dotenv import load_dotenv

# load secrets
load_dotenv()
SD_KEY = os.getenv("SD_KEY")
HF_KEY = os.getenv("HF_KEY")
USERNAME = os.getenv("USERNAME")
PASS = os.getenv("PASS")

# set up dataset writer
hf_writer_gen = gr.HuggingFaceDatasetSaver(HF_KEY, "helms/master-thesis-generated-images", private=True, separate_dirs=False)

def infer(prompt):
    stability_api = client.StabilityInference(
        key=SD_KEY,  # AaPI Key reference.
        verbose=True,  # Print debug messages.
        engine="stable-diffusion-xl-1024-v1-0",  # Set the engine to use for generation.
        # Available engines: stable-diffusion-xl-1024-v1-0 stable-diffusion-xl-1024-v0-9 stable-diffusion-v1 stable-diffusion-v1-5 stable-diffusion-512-v2-0 stable-diffusion-768-v2-0
        # stable-diffusion-512-v2-1 stable-diffusion-768-v2-1 stable-diffusion-xl-beta-v2-2-2 stable-inpainting-v1-0 stable-inpainting-512-v2-0
    )
    answers = stability_api.generate(
        prompt=prompt,
        # seed=992446758, # If a seed is provided, the resulting generated image will be deterministic.
        # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again.
        # Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook.
        steps=30,  # Amount of inference steps performed on image generation. Defaults to 30.
        cfg_scale=7.0,  # Influences how strongly your generation is guided to match your prompt.
        # Setting this value higher increases the strength in which it tries to match your prompt.
        # Defaults to 7.0 if not specified.
        width=1024,  # Generation width, defaults to 512 if not included.
        height=1024,  # Generation height, defaults to 512 if not included.
        samples=1,  # Number of images to generate, defaults to 1 if not included.
        sampler=generation.SAMPLER_K_DPMPP_2M  # Choose which sampler we want to denoise our generation with.
        # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers.
        # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m)
    )
    for resp in answers:
        for artifact in resp.artifacts:
            if artifact.finish_reason == generation.FILTER:
                warnings.warn(
                    "Your request activated the APIs safety filters and could not be processed."
                    "Please modify the prompt and try again.")
            if artifact.type == generation.ARTIFACT_IMAGE:
                img = Image.open(io.BytesIO(artifact.binary))
                return img


with gr.Blocks(title="Master Thesis Image Generator") as demo:
    with gr.Row():
        gr.HTML(""" 
                    <div style="text-align: justify; margin: 0 auto;">
                    <h1><a id="Master_Thesis_Image_Generator_0"></a>Master Thesis Image Generator</h1>
                    <p>Welcome to this demo app for the Master thesis of Fabian Helms, TU Dortmund.<br>
                    Contact: <a href="mailto:fabian.helms@tu-dortmund.de">fabian.helms@tu-dortmund.de</a></p>
                    <li>Please enter your name on the left hand side.</li>
                    <li>In the prompt field you can instruct the model on what you want to generate.</li>
                    <li>Generated images will be stored for research purposes only.</li>
                    </div>
                """)
        gr.HTML("""
                    <div style="text-align: justify; margin: 0 auto;">
                    <br>Details:
                    <li>English prompts yield best results. Examples are shown at the bottom.</li>
                    <li>The generation of an image takes about 20-30 seconds.</li>
                    <li>Generated images are 1024x1024px in size.</li>
                    <li>You can see your history of generations on the left hand side. Click on an image to view it in full size.</li>
                    <li>Model: <a href="https://stability.ai/blog/stable-diffusion-sdxl-1-announcement">Stable Diffusion XL 1.0</a></li>
                    </div>
                """)
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            user = gr.Textbox(placeholder="Please enter your name.", label="Name")
            inp = gr.Textbox(placeholder="Enter your prompt here.", lines=5, label="Prompt")
            btn_gen = gr.Button("Generate Image", variant="primary", size="lg", visible=True)
            with gr.Box():
                gr.HTML(
                    """<div style="text-align: center; margin: 0 auto; padding-bottom: 10px"><p style="font-size:10pt">Generated Images</p></div>""")
                store = gr.Gallery(value=[], visible=True, show_share_button=False, allow_preview=False, columns=1)
        with gr.Column(scale=3):
            out = gr.Image(interactive=False, visible=True, container=True, width=1024, show_share_button=False)
    with gr.Row():
        with gr.Accordion("See Examples:", open=False):
            examples = [
                ["Vintage hot rod with custom flame paint job"],
                ["Ancient, mysterious temple in a mountain range, surrounded by misty clouds and tall peaks"],
                ["Glimpses of a herd of wild elephants crossing a savanna"],
                ["Beautiful waterfall in a lush jungle, with sunlight shining through the trees"]
            ]
            ex = gr.Examples(examples=examples, inputs=inp)

    def make_visible():
        return gr.update(visible=True)

    def make_invisible():
        return gr.update(visible=False)

    def make_inactive():
        return gr.update(interactive=False)

    def make_active():
        return gr.update(interactive=True)

    def add_to_gallery(img, gall):
        return [img] + [d['name'] for d in gall]

    def select_img(evt: gr.SelectData, gall):
        sel_img = [d['name'] for d in gall][evt.index]
        return sel_img

    hf_writer_gen.setup([user, inp, out], ".temp_gen")

    btn_gen.click(
        fn=make_inactive, outputs=btn_gen).then(
        fn=infer, inputs=inp, outputs=out, scroll_to_output=True).then(
        fn=add_to_gallery, inputs=[out, store], outputs=store).then(
        fn=make_active, outputs=btn_gen).then(
        lambda *args: hf_writer_gen.flag(args), inputs=[user, inp, out], outputs=None, preprocess=False)

    inp.submit(
        fn=make_inactive, outputs=btn_gen).then(
        fn=infer, inputs=inp, outputs=out, scroll_to_output=True).then(
        fn=add_to_gallery, inputs=[out, store], outputs=store).then(
        fn=make_active, outputs=btn_gen).then(
        lambda *args: hf_writer_gen.flag(args), inputs=[user, inp, out], outputs=None, preprocess=False)

    store.select(fn=select_img, inputs=store, outputs=out)

demo.launch(show_api=False, auth=(USERNAME, PASS)
           )