Create app.py

app.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import math
+import random
+
+import gradio as gr
+import torch
+from PIL import Image, ImageOps
+from diffusers import StableDiffusionSAGPipeline
+
+
+help_text = """
+Self-Attention Guidance
+"""
+
+
+examples = [
+    [
+        ' ',
+        50,
+        False,
+        8978,
+        7.5,
+        1.0,
+    ],
+    [
+        '.',
+        50,
+        False,
+        8978,
+        7.5,
+        1.0,
+    ],
+    [
+        'A cute Scottish Fold playing with a ball',
+        50,
+        False,
+        8978,
+        5.0,
+        1.0,
+    ],
+    [
+        'A person with a happy dog',
+        50,
+        False,
+        8978,
+        5.0,
+        1.0,
+    ],
+]
+
+
+model_id = "runwayml/stable-diffusion-v1-5"
+
+def main():
+    pipe = StableDiffusionSAGPipeline.from_pretrained(model_id)#, torch_dtype=torch.float16)
+
+    def generate(
+        prompt: str,
+        steps: int,
+        randomize_seed: bool,
+        seed: int,
+        cfg_scale: float,
+        sag_scale: float,
+    ):
+        seed = random.randint(0, 100000) if randomize_seed else seed
+
+        generator = torch.manual_seed(seed)
+        ori_image = pipe(prompt, generator=generator, guidance_scale=cfg_scale, sag_scale=0.75).images[0]
+        generator = torch.manual_seed(seed)
+        sag_image = pipe(prompt, generator=generator, guidance_scale=cfg_scale, sag_scale=0.75).images[0]
+        return [seed, ori_image, sag_image]
+
+    def reset():
+        return [0, "Randomize Seed", 1371, 5.0, 0.75, None, None]
+
+    with gr.Blocks() as demo:
+        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">
+           Self-Attention Guidance
+        """)
+        with gr.Row():
+            with gr.Column(scale=5):
+                prompt = gr.Textbox(lines=1, label="Enter your prompt", interactive=True)
+            with gr.Column(scale=1, min_width=60):
+                generate_button = gr.Button("Generate")
+            with gr.Column(scale=1, min_width=60):
+                reset_button = gr.Button("Reset")
+
+        with gr.Row():
+            ori_image = gr.Image(label="CFG", type="pil", interactive=False)
+            sag_image = gr.Image(label="SAG + CFG", type="pil", interactive=False)
+            ori_image.style(height=512, width=512)
+            sag_image.style(height=512, width=512)
+
+        with gr.Row():
+            steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)
+            randomize_seed = gr.Radio(
+                ["Fix Seed", "Randomize Seed"],
+                value="Fix Seed",
+                type="index",
+                show_label=False,
+                interactive=True,
+            )
+            seed = gr.Number(value=8978, precision=0, label="Seed", interactive=True)
+            
+        with gr.Row():
+            cfg_scale = gr.Slider(
+                label="Guidance Scale", minimum=0, maximum=10, value=5.0, step=0.1
+            )
+            sag_scale = gr.Slider(
+                label="Self-Attention Guidance Scale", minimum=0, maximum=1.0, value=0.75, step=0.05
+            )
+            
+        ex = gr.Examples(
+            examples=examples,
+            fn=generate,
+            inputs=[
+                prompt,
+                steps,
+                randomize_seed,
+                seed,
+                cfg_scale,
+                sag_scale,
+            ],
+            outputs=[seed, ori_image, sag_image],
+            cache_examples=True,
+            preprocess=False,
+            postprocess=False
+        )
+
+        gr.Markdown(help_text)
+        
+        generate_button.click(
+            fn=generate,
+            inputs=[
+                prompt,
+                steps,
+                randomize_seed,
+                seed,
+                cfg_scale,
+                sag_scale,
+            ],
+            outputs=[seed, ori_image, sag_image],
+        )
+        reset_button.click(
+            fn=reset,
+            inputs=[],
+            outputs=[steps, randomize_seed, seed, cfg_scale, sag_scale, ori_image, sag_image],
+        )
+
+    demo.queue(concurrency_count=1)
+    demo.launch(share=False)
+
+
+if __name__ == "__main__":
+    main()