Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 
 import torch
 import torchvision.transforms.functional as F
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, AutoencoderTiny, AutoencoderKL
 import gradio as gr
 
 device = "cuda"
@@ -19,6 +19,16 @@ pipe = StableDiffusionControlNetPipeline.from_pretrained(
 )
 pipe.to(device)
 
+vae_tiny = AutoencoderTiny.from_pretrained(
+    "IDKiro/sdxs-512-dreamshaper", subfolder="vae"
+)
+vae_tiny.to(device, dtype=weight_type)
+
+vae_large = AutoencoderKL.from_pretrained(
+    "IDKiro/sdxs-512-dreamshaper", subfolder="vae_large"
+)
+vae_tiny.to(device, dtype=weight_type)
+
 style_list = [
     {
         "name": "No Style",
@@ -81,9 +91,15 @@ def run(
     prompt_template,
     style_name,
     controlnet_conditioning_scale,
+    vae_type="tiny vae",
     device_type="GPU",
     param_dtype="torch.float16",
 ):
+    if vae_type == "tiny vae":
+        pipe.vae = vae_tiny
+    elif vae_type == "large vae":
+        pipe.vae = vae_large
+
     if device_type == "CPU":
         device = "cpu"
         param_dtype = "torch.float32"
@@ -118,24 +134,28 @@ def run(
     return output_pil
 
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme="monochrome") as demo:
     gr.Markdown("# SDXS-512-DreamShaper-Sketch")
-    gr.Markdown("[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)")
-    with gr.Row(elem_id="main_row"):
-        with gr.Column(elem_id="column_input"):
-            gr.Markdown("## INPUT", elem_id="input_header")
+    gr.Markdown(
+        "[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)"
+    )
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## INPUT")
             image = gr.Sketchpad(
                 type="pil",
                 image_mode="RGBA",
                 brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=8),
-                crop_size=(512, 512),
+                crop_size="1:1",
             )
 
-            # gr.Markdown("## Prompt", elem_id="tools_header")
             prompt = gr.Textbox(label="Prompt", value="", show_label=True)
             with gr.Row():
                 style = gr.Dropdown(
-                    label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME, scale=1
+                    label="Style",
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                    scale=1,
                 )
                 prompt_temp = gr.Textbox(
                     label="Prompt Style Template",
@@ -148,6 +168,15 @@ with gr.Blocks() as demo:
                 label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.8
             )
 
+            vae_choices = ["tiny vae", "large vae"]
+            vae_type = gr.Radio(
+                vae_choices,
+                label="Image Decoder Type",
+                value=vae_choices[0],
+                interactive=True,
+                info="To save GPU memory, use tiny vae. For better quality, use large vae.",
+            )
+
             device_choices = ["GPU", "CPU"]
             device_type = gr.Radio(
                 device_choices,
@@ -166,16 +195,19 @@ with gr.Blocks() as demo:
                 info="To save GPU memory, use torch.float16. For better quality, use torch.float32.",
             )
 
-        with gr.Column(elem_id="column_output"):
-            gr.Markdown("## OUTPUT", elem_id="output_header")
+        with gr.Column():
+            gr.Markdown("## OUTPUT")
             result = gr.Image(
                 label="Result",
-                height=512,
-                width=512,
-                elem_id="output_image",
                 show_label=False,
                 show_download_button=True,
             )
+            run_button = gr.Button("Run")
+            gr.Markdown("### Instructions")
+            gr.Markdown("**1**. Enter a text prompt (e.g. cat)")
+            gr.Markdown("**2**. Start sketching")
+            gr.Markdown("**3**. Change the image style using a style template")
+            gr.Markdown("**4**. Adjust the effect of sketch guidance using the slider")
 
     inputs = [
         image,
@@ -183,6 +215,7 @@ with gr.Blocks() as demo:
         prompt_temp,
         style,
         controlnet_conditioning_scale,
+        vae_type,
         device_type,
         param_dtype,
     ]
@@ -190,9 +223,25 @@ with gr.Blocks() as demo:
 
     prompt.change(fn=run, inputs=inputs, outputs=outputs)
     style.change(lambda x: styles[x], inputs=[style], outputs=[prompt_temp]).then(
-            fn=run, inputs=inputs, outputs=outputs,)
-    image.change(run, inputs=inputs, outputs=outputs,)
-    controlnet_conditioning_scale.change(run, inputs=inputs, outputs=outputs,)
+        fn=run,
+        inputs=inputs,
+        outputs=outputs,
+    )
+    image.change(
+        run,
+        inputs=inputs,
+        outputs=outputs,
+    )
+    controlnet_conditioning_scale.change(
+        run,
+        inputs=inputs,
+        outputs=outputs,
+    )
+    run_button.click(
+        run,
+        inputs=inputs,
+        outputs=outputs,
+    )
 
 if __name__ == "__main__":
     demo.queue().launch()