Migrate from yapf to black

.pre-commit-config.yaml

@@ -1,7 +1,6 @@
-exclude: patch
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.2.0
+  rev: v4.4.0
   hooks:
   - id: check-executables-have-shebangs
   - id: check-json
@@ -9,29 +8,43 @@ repos:
   - id: check-shebang-scripts-are-executable
   - id: check-toml
   - id: check-yaml
-  - id: double-quote-string-fixer
   - id: end-of-file-fixer
   - id: mixed-line-ending
-    args: ['--fix=lf']
+    args: ["--fix=lf"]
   - id: requirements-txt-fixer
   - id: trailing-whitespace
 - repo: https://github.com/myint/docformatter
-  rev: v1.4
+  rev: v1.7.5
   hooks:
   - id: docformatter
-    args: ['--in-place']
+    args: ["--in-place"]
 - repo: https://github.com/pycqa/isort
   rev: 5.12.0
   hooks:
     - id: isort
+      args: ["--profile", "black"]
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.991
+  rev: v1.5.1
   hooks:
     - id: mypy
-      args: ['--ignore-missing-imports']
-      additional_dependencies: ['types-python-slugify']
-- repo: https://github.com/google/yapf
-  rev: v0.32.0
+      args: ["--ignore-missing-imports"]
+      additional_dependencies: ["types-python-slugify", "types-requests", "types-PyYAML"]
+- repo: https://github.com/psf/black
+  rev: 23.9.1
   hooks:
-  - id: yapf
-    args: ['--parallel', '--in-place']
+    - id: black
+      language_version: python3.10
+      args: ["--line-length", "119"]
+- repo: https://github.com/kynan/nbstripout
+  rev: 0.6.1
+  hooks:
+    - id: nbstripout
+      args: ["--extra-keys", "metadata.interpreter metadata.kernelspec cell.metadata.pycharm"]
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.7.0
+  hooks:
+    - id: nbqa-black
+    - id: nbqa-pyupgrade
+      args: ["--py37-plus"]
+    - id: nbqa-isort
+      args: ["--float-to-top"]

.style.yapf

@@ -1,5 +0,0 @@
-[style]
-based_on_style = pep8
-blank_line_before_nested_class_or_def = false
-spaces_before_comment = 2
-split_before_logical_operator = true

.vscode/settings.json

@@ -1,18 +1,21 @@
 {
-    "python.linting.enabled": true,
-    "python.linting.flake8Enabled": true,
-    "python.linting.pylintEnabled": false,
-    "python.linting.lintOnSave": true,
-    "python.formatting.provider": "yapf",
-    "python.formatting.yapfArgs": [
-        "--style={based_on_style: pep8, indent_width: 4, blank_line_before_nested_class_or_def: false, spaces_before_comment: 2, split_before_logical_operator: true}"
-    ],
     "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter",
         "editor.formatOnType": true,
         "editor.codeActionsOnSave": {
             "source.organizeImports": true
         }
     },
+    "black-formatter.args": [
+        "--line-length=119"
+    ],
+    "isort.args": ["--profile", "black"],
+    "flake8.args": [
+        "--max-line-length=119"
+    ],
+    "ruff.args": [
+        "--line-length=119"
+    ],
     "editor.formatOnSave": true,
     "files.insertFinalNewline": true
 }

app.py

@@ -9,11 +9,13 @@ import PIL.Image
 
 from model import Model
 
-DESCRIPTION = '''# Attend-and-Excite
+DESCRIPTION = """\
+# Attend-and-Excite
+
 This is a demo for [Attend-and-Excite](https://arxiv.org/abs/2301.13826).
 Attend-and-Excite performs attention-based generative semantic guidance to mitigate subject neglect in Stable Diffusion.
 Select a prompt and a set of indices matching the subjects you wish to strengthen (the `Check token indices` cell can help map between a word and its index).
-'''
+"""
 
 model = Model()
 
@@ -28,148 +30,157 @@ def process_example(
     guidance_scale = 7.5
 
     token_table = model.get_token_table(prompt)
-    result = model.run(prompt, indices_to_alter_str, seed,
-                       apply_attend_and_excite, num_steps, guidance_scale)
+    result = model.run(prompt, indices_to_alter_str, seed, apply_attend_and_excite, num_steps, guidance_scale)
     return token_table, result
 
 
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
 
     with gr.Row():
         with gr.Column():
             prompt = gr.Text(
-                label='Prompt',
+                label="Prompt",
                 max_lines=1,
-                placeholder=
-                'A pod of dolphins leaping out of the water in an ocean with a ship on the background'
+                placeholder="A pod of dolphins leaping out of the water in an ocean with a ship on the background",
             )
-            with gr.Accordion(label='Check token indices', open=False):
-                show_token_indices_button = gr.Button('Show token indices')
-                token_indices_table = gr.Dataframe(label='Token indices',
-                                                   headers=['Index', 'Token'],
-                                                   col_count=2)
+            with gr.Accordion(label="Check token indices", open=False):
+                show_token_indices_button = gr.Button("Show token indices")
+                token_indices_table = gr.Dataframe(label="Token indices", headers=["Index", "Token"], col_count=2)
             token_indices_str = gr.Text(
-                label=
-                'Token indices (a comma-separated list indices of the tokens you wish to alter)',
+                label="Token indices (a comma-separated list indices of the tokens you wish to alter)",
                 max_lines=1,
-                placeholder='4,16')
-            seed = gr.Slider(label='Seed',
-                             minimum=0,
-                             maximum=100000,
-                             value=0,
-                             step=1)
-            apply_attend_and_excite = gr.Checkbox(
-                label='Apply Attend-and-Excite', value=True)
-            num_steps = gr.Slider(label='Number of steps',
-                                  minimum=0,
-                                  maximum=100,
-                                  step=1,
-                                  value=50)
-            guidance_scale = gr.Slider(label='CFG scale',
-                                       minimum=0,
-                                       maximum=50,
-                                       step=0.1,
-                                       value=7.5)
-            run_button = gr.Button('Generate')
+                placeholder="4,16",
+            )
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=100000,
+                step=1,
+                value=0,
+            )
+            apply_attend_and_excite = gr.Checkbox(label="Apply Attend-and-Excite", value=True)
+            num_steps = gr.Slider(
+                label="Number of steps",
+                minimum=0,
+                maximum=100,
+                step=1,
+                value=50,
+            )
+            guidance_scale = gr.Slider(
+                label="CFG scale",
+                minimum=0,
+                maximum=50,
+                step=0.1,
+                value=7.5,
+            )
+            run_button = gr.Button("Generate")
         with gr.Column():
-            result = gr.Image(label='Result')
+            result = gr.Image(label="Result")
 
     with gr.Row():
         examples = [
             [
-                'A mouse and a red car',
-                '2,6',
+                "A mouse and a red car",
+                "2,6",
                 2098,
                 True,
             ],
             [
-                'A mouse and a red car',
-                '2,6',
+                "A mouse and a red car",
+                "2,6",
                 2098,
                 False,
             ],
             [
-                'A horse and a dog',
-                '2,5',
+                "A horse and a dog",
+                "2,5",
                 123,
                 True,
             ],
             [
-                'A horse and a dog',
-                '2,5',
+                "A horse and a dog",
+                "2,5",
                 123,
                 False,
             ],
             [
-                'A painting of an elephant with glasses',
-                '5,7',
+                "A painting of an elephant with glasses",
+                "5,7",
                 123,
                 True,
             ],
             [
-                'A painting of an elephant with glasses',
-                '5,7',
+                "A painting of an elephant with glasses",
+                "5,7",
                 123,
                 False,
             ],
             [
-                'A playful kitten chasing a butterfly in a wildflower meadow',
-                '3,6,10',
+                "A playful kitten chasing a butterfly in a wildflower meadow",
+                "3,6,10",
                 123,
                 True,
             ],
             [
-                'A playful kitten chasing a butterfly in a wildflower meadow',
-                '3,6,10',
+                "A playful kitten chasing a butterfly in a wildflower meadow",
+                "3,6,10",
                 123,
                 False,
             ],
             [
-                'A grizzly bear catching a salmon in a crystal clear river surrounded by a forest',
-                '2,6,15',
+                "A grizzly bear catching a salmon in a crystal clear river surrounded by a forest",
+                "2,6,15",
                 123,
                 True,
             ],
             [
-                'A grizzly bear catching a salmon in a crystal clear river surrounded by a forest',
-                '2,6,15',
+                "A grizzly bear catching a salmon in a crystal clear river surrounded by a forest",
+                "2,6,15",
                 123,
                 False,
             ],
             [
-                'A pod of dolphins leaping out of the water in an ocean with a ship on the background',
-                '4,16',
+                "A pod of dolphins leaping out of the water in an ocean with a ship on the background",
+                "4,16",
                 123,
                 True,
             ],
             [
-                'A pod of dolphins leaping out of the water in an ocean with a ship on the background',
-                '4,16',
+                "A pod of dolphins leaping out of the water in an ocean with a ship on the background",
+                "4,16",
                 123,
                 False,
             ],
         ]
-        gr.Examples(examples=examples,
-                    inputs=[
-                        prompt,
-                        token_indices_str,
-                        seed,
-                        apply_attend_and_excite,
-                    ],
-                    outputs=[
-                        token_indices_table,
-                        result,
-                    ],
-                    fn=process_example,
-                    cache_examples=os.getenv('CACHE_EXAMPLES') == '1',
-                    examples_per_page=20)
+        gr.Examples(
+            examples=examples,
+            inputs=[
+                prompt,
+                token_indices_str,
+                seed,
+                apply_attend_and_excite,
+            ],
+            outputs=[
+                token_indices_table,
+                result,
+            ],
+            fn=process_example,
+            cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
+            examples_per_page=20,
+        )
 
     show_token_indices_button.click(
         fn=model.get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
+        api_name=False,
     )
 
     inputs = [
@@ -185,31 +196,37 @@ with gr.Blocks(css='style.css') as demo:
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
+        api_name=False,
     ).then(
         fn=model.run,
         inputs=inputs,
         outputs=result,
+        api_name=False,
     )
     token_indices_str.submit(
         fn=model.get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
+        api_name=False,
     ).then(
         fn=model.run,
         inputs=inputs,
         outputs=result,
+        api_name=False,
     )
     run_button.click(
         fn=model.get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
+        api_name=False,
     ).then(
         fn=model.run,
         inputs=inputs,
         outputs=result,
-        api_name='run',
+        api_name="run",
     )
 
-demo.queue(max_size=10).launch()
+if __name__ == "__main__":
+    demo.queue(max_size=10).launch()

model.py

@@ -2,26 +2,20 @@ from __future__ import annotations
 
 import PIL.Image
 import torch
-from diffusers import (StableDiffusionAttendAndExcitePipeline,
-                       StableDiffusionPipeline)
+from diffusers import StableDiffusionAttendAndExcitePipeline, StableDiffusionPipeline
 
 
 class Model:
     def __init__(self):
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
-        model_id = 'CompVis/stable-diffusion-v1-4'
-        self.ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(
-            model_id)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        model_id = "CompVis/stable-diffusion-v1-4"
+        self.ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(model_id)
         self.ax_pipe.to(self.device)
         self.sd_pipe = StableDiffusionPipeline.from_pretrained(model_id)
         self.sd_pipe.to(self.device)
 
     def get_token_table(self, prompt: str):
-        tokens = [
-            self.ax_pipe.tokenizer.decode(t)
-            for t in self.ax_pipe.tokenizer(prompt)['input_ids']
-        ]
+        tokens = [self.ax_pipe.tokenizer.decode(t) for t in self.ax_pipe.tokenizer(prompt)["input_ids"]]
         tokens = tokens[1:-1]
         return list(enumerate(tokens, start=1))
 
@@ -44,9 +38,9 @@ class Model:
 
         if apply_attend_and_excite:
             try:
-                token_indices = list(map(int, indices_to_alter_str.split(',')))
+                token_indices = list(map(int, indices_to_alter_str.split(",")))
             except Exception:
-                raise ValueError('Invalid token indices.')
+                raise ValueError("Invalid token indices.")
             out = self.ax_pipe(
                 prompt=prompt,
                 token_indices=token_indices,

style.css

@@ -1,3 +1,10 @@
 h1 {
   text-align: center;
 }
+
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+}