Add application file

app.py

@@ -0,0 +1,153 @@
+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
+
+
+# theme = gr.themes.Monochrome(
+#     primary_hue="indigo",
+#     secondary_hue="blue",
+#     neutral_hue="slate",
+#     radius_size=gr.themes.sizes.radius_sm,
+#     font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
+# )
+load_dotenv()
+SD_KEY = os.getenv("SD_KEY")
+HF_KEY = os.getenv("HF_KEY")
+USERNAME = os.getenv("USERNAME")
+PASS = os.getenv("PASS")
+
+
+hf_writer_gen = gr.HuggingFaceDatasetSaver(HF_KEY, "helms/master-thesis-generated-images", private=True, separate_dirs=False)
+hf_writer_up = gr.HuggingFaceDatasetSaver(HF_KEY, "helms/master-thesis-upscaled-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-v1",  # Set the engine to use for generation.
+        # Available engines: stable-diffusion-v1 stable-diffusion-v1-5 stable-diffusion-512-v2-0 stable-diffusion-768-v2-0 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=512,  # Generation width, defaults to 512 if not included.
+        height=512,  # Generation height, defaults to 512 if not included.
+        samples=4,  # 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)
+    )
+    infer.results = []
+    for resp in answers:
+        for artifact in resp.artifacts:
+            if artifact.finish_reason == generation.FILTER:
+                warnings.warn(
+                    "Your request activated the API's 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))
+                infer.results.append(img)
+                #img.save(fp=str(prompt) + ".png")
+
+    return infer.results
+
+def upscale(image):
+    # Set up our connection to the API.
+    stability_api = client.StabilityInference(
+        key=SD_KEY,  # API Key reference.
+        upscale_engine="esrgan-v1-x2plus",  # The name of the upscaling model we want to use.
+        # Available Upscaling Engines: esrgan-v1-x2plus, stable-diffusion-x4-latent-upscaler
+        verbose=True,  # Print debug messages.
+    )
+    # Import our local image to use as a reference for our upscaled image.
+    # The 'img' variable below is set to a local file for upscaling, however if you are already running a generation call and have an image artifact available, you can pass that image artifact to the upscale function instead.
+    img = Image.open(image)
+
+    answers = stability_api.upscale(
+        init_image=img,  # Pass our image to the API and call the upscaling process.
+        width=1024, # Optional parameter to specify the desired output width.
+        # prompt="A beautiful sunset", # Optional parameter when using `stable-diffusion-x4-latent-upscaler` to specify a prompt to use for the upscaling process.
+        # seed=1234, # Optional parameter when using `stable-diffusion-x4-latent-upscaler` to specify a seed to use for the upscaling process.
+        # steps=20, # Optional parameter when using `stable-diffusion-x4-latent-upscaler` to specify the number of diffusion steps to use for the upscaling process. Defaults to 20 if no value is passed, with a maximum of 50.
+        # cfg_scale=7 # Optional parameter when using `stable-diffusion-x4-latent-upscaler` to specify the strength of prompt in use for the upscaling process. Defaults to 7 if no value is passed.
+    )
+    # Set up our warning to print to the console if the adult content classifier is tripped.
+    # If adult content classifier is not tripped, save our image.
+
+    for resp in answers:
+        for artifact in resp.artifacts:
+            if artifact.finish_reason == generation.FILTER:
+                warnings.warn(
+                    "Your request activated the API's safety filters and could not be processed."
+                    "Please submit a different image and try again.")
+            if artifact.type == generation.ARTIFACT_IMAGE:
+                big_img = Image.open(io.BytesIO(artifact.binary))
+                #big_img.save("imageupscaled" + ".png")  # Save our image to a local file.
+
+    return big_img
+
+
+# demo = gr.Interface(
+#         infer,
+#         gr.Textbox(label="Input", lines=5),
+#         gr.Gallery(label="Output")
+# )
+
+with gr.Blocks(title="Master Thesis Image Generator") as demo:
+    gr.Markdown("Welcome to this demo app for the Master thesis of Fabian Helms, TU Dortmund")
+    with gr.Row():
+        with gr.Column():
+            user = gr.Textbox(placeholder="Please enter your name.", label="Name")
+            inp = gr.Textbox(placeholder="Enter your prompt here.", lines=3, label="Prompt")
+            btn_gen = gr.Button("Generate Images", variant="primary", size="lg")
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            out_gen = gr.Gallery(label="Generated Images", columns=2, rows=2, container=True, preview=False, height="2", allow_preview=True)
+            selected = gr.Image(label="selected img", type="filepath", interactive=False, visible=False)#.style(height=500, width=500)
+        with gr.Column(scale=1):
+            with gr.Row():
+                gr.Markdown(value="Please select the image you want to upscale.", visible=True)
+                btn_upscale = gr.Button("Upscale selected image", variant="secondary", size="sm")
+            with gr.Row():
+                out_upscale = gr.Image(label="Upscaled Image", interactive=False)#.style(height=500, width=500)
+                #btn_flag = gr.Button("Flag Image", variant="secondary", size="sm")
+
+    hf_writer_gen.setup([user, inp, out_gen], ".temp_gen")
+    hf_writer_up.setup([user, inp, out_upscale], ".temp_upscale")
+    def select_img(evt: gr.SelectData):
+        sel_img = infer.results[evt.index]
+        #hf_writer.flag([inp, user, out_gen])
+        #sel_img.save("check.png")
+        return sel_img
+
+    def gen(prompt):
+        res_gen = infer(prompt)
+        return out_gen.update(value=res_gen)
+
+
+    inp.submit(fn=infer, inputs=inp, outputs=out_gen
+               ).then(lambda *args: hf_writer_gen.flag(args), inputs=[user, inp, out_gen], outputs=None, preprocess=False)
+    btn_gen.click(fn=infer, inputs=inp, outputs=out_gen
+                ).then(lambda *args: hf_writer_gen.flag(args), inputs=[user, inp, out_gen], outputs=None, preprocess=False)
+
+    btn_upscale.click(fn=upscale, inputs=selected, outputs=out_upscale
+                      ).then(lambda *args: hf_writer_up.flag(args), inputs=[user, inp, out_upscale], outputs=None, preprocess=False)
+
+    out_gen.select(select_img, outputs=selected)
+    #btn_flag.click(lambda *args: hf_writer.flag(args), inputs=[user, inp, out_gen], outputs=None, preprocess=False)
+
+
+
+demo.launch(show_api=False, auth=(USERNAME, PASS))