Migrate from yapf to black

.pre-commit-config.yaml

@@ -1,36 +1,55 @@
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.2.0
-  hooks:
-  - id: check-executables-have-shebangs
-  - id: check-json
-  - id: check-merge-conflict
-  - 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']
-  - id: requirements-txt-fixer
-  - id: trailing-whitespace
-- repo: https://github.com/myint/docformatter
-  rev: v1.4
-  hooks:
-  - id: docformatter
-    args: ['--in-place']
-- repo: https://github.com/pycqa/isort
-  rev: 5.12.0
-  hooks:
-    - id: isort
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.991
-  hooks:
-    - id: mypy
-      args: ['--ignore-missing-imports']
-      additional_dependencies: ['types-python-slugify']
-- repo: https://github.com/google/yapf
-  rev: v0.32.0
-  hooks:
-  - id: yapf
-    args: ['--parallel', '--in-place']
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: check-executables-have-shebangs
+      - id: check-json
+      - id: check-merge-conflict
+      - id: check-shebang-scripts-are-executable
+      - id: check-toml
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+      - id: requirements-txt-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/myint/docformatter
+    rev: v1.7.5
+    hooks:
+      - id: docformatter
+        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: v1.5.1
+    hooks:
+      - id: mypy
+        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: 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

@@ -11,13 +11,13 @@ import torch
 
 from model import Model
 
-DESCRIPTION = '# [UniDiffuser](https://github.com/thu-ml/unidiffuser)'
+DESCRIPTION = "# [UniDiffuser](https://github.com/thu-ml/unidiffuser)"
 
-SPACE_ID = os.getenv('SPACE_ID')
+SPACE_ID = os.getenv("SPACE_ID")
 if SPACE_ID is not None:
     DESCRIPTION += f'\n<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'
 if not torch.cuda.is_available():
-    DESCRIPTION += '\n<p>Running on CPU 🥶</p>'
+    DESCRIPTION += "\n<p>Running on CPU 🥶</p>"
 
 model = Model()
 
@@ -34,50 +34,49 @@ def create_demo(mode_name: str) -> gr.Blocks:
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
-                mode = gr.Dropdown(label='Mode',
-                                   choices=[
-                                       't2i',
-                                       'i2t',
-                                       'joint',
-                                       'i',
-                                       't',
-                                       'i2t2i',
-                                       't2i2t',
-                                   ],
-                                   value=mode_name,
-                                   visible=False)
-                prompt = gr.Text(label='Prompt',
-                                 max_lines=1,
-                                 visible=mode_name in ['t2i', 't2i2t'])
-                image = gr.Image(label='Input image',
-                                 type='pil',
-                                 visible=mode_name in ['i2t', 'i2t2i'])
-                run_button = gr.Button('Run')
-                with gr.Accordion('Advanced options', open=False):
-                    seed = gr.Slider(label='Seed',
-                                     minimum=0,
-                                     maximum=MAX_SEED,
-                                     step=1,
-                                     value=0)
-                    randomize_seed = gr.Checkbox(label='Randomize seed',
-                                                 value=True)
-                    num_steps = gr.Slider(label='Steps',
-                                          minimum=1,
-                                          maximum=100,
-                                          value=20,
-                                          step=1)
-                    guidance_scale = gr.Slider(label='Guidance Scale',
-                                               minimum=0.1,
-                                               maximum=30.0,
-                                               value=8.0,
-                                               step=0.1)
+                mode = gr.Dropdown(
+                    label="Mode",
+                    choices=[
+                        "t2i",
+                        "i2t",
+                        "joint",
+                        "i",
+                        "t",
+                        "i2t2i",
+                        "t2i2t",
+                    ],
+                    value=mode_name,
+                    visible=False,
+                )
+                prompt = gr.Text(label="Prompt", max_lines=1, visible=mode_name in ["t2i", "t2i2t"])
+                image = gr.Image(label="Input image", type="pil", visible=mode_name in ["i2t", "i2t2i"])
+                run_button = gr.Button("Run")
+                with gr.Accordion("Advanced options", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                    num_steps = gr.Slider(
+                        label="Steps",
+                        minimum=1,
+                        maximum=100,
+                        value=20,
+                        step=1,
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=0.1,
+                        maximum=30.0,
+                        value=8.0,
+                        step=0.1,
+                    )
             with gr.Column():
-                result_image = gr.Image(label='Generated image',
-                                        visible=mode_name
-                                        in ['t2i', 'i', 'joint', 'i2t2i'])
-                result_text = gr.Text(label='Generated text',
-                                      visible=mode_name
-                                      in ['i2t', 't', 'joint', 't2i2t'])
+                result_image = gr.Image(label="Generated image", visible=mode_name in ["t2i", "i", "joint", "i2t2i"])
+                result_text = gr.Text(label="Generated text", visible=mode_name in ["i2t", "t", "joint", "t2i2t"])
         inputs = [
             mode,
             prompt,
@@ -110,26 +109,28 @@ def create_demo(mode_name: str) -> gr.Blocks:
             fn=model.run,
             inputs=inputs,
             outputs=outputs,
-            api_name=f'run_{mode_name}',
+            api_name=f"run_{mode_name}",
         )
     return demo
 
 
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Tabs():
-        with gr.TabItem('text2image'):
-            create_demo('t2i')
-        with gr.TabItem('image2text'):
-            create_demo('i2t')
-        with gr.TabItem('image variation'):
-            create_demo('i2t2i')
-        with gr.TabItem('joint generation'):
-            create_demo('joint')
-        with gr.TabItem('image generation'):
-            create_demo('i')
-        with gr.TabItem('text generation'):
-            create_demo('t')
-        with gr.TabItem('text variation'):
-            create_demo('t2i2t')
-demo.queue(max_size=15).launch()
+        with gr.TabItem("text2image"):
+            create_demo("t2i")
+        with gr.TabItem("image2text"):
+            create_demo("i2t")
+        with gr.TabItem("image variation"):
+            create_demo("i2t2i")
+        with gr.TabItem("joint generation"):
+            create_demo("joint")
+        with gr.TabItem("image generation"):
+            create_demo("i")
+        with gr.TabItem("text generation"):
+            create_demo("t")
+        with gr.TabItem("text variation"):
+            create_demo("t2i2t")
+
+if __name__ == "__main__":
+    demo.queue(max_size=15).launch()

model.py

@@ -7,15 +7,12 @@ from diffusers import UniDiffuserPipeline
 
 class Model:
     def __init__(self):
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
-        if self.device.type == 'cuda':
-            self.pipe = UniDiffuserPipeline.from_pretrained(
-                'thu-ml/unidiffuser-v1', torch_dtype=torch.float16)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        if self.device.type == "cuda":
+            self.pipe = UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1", torch_dtype=torch.float16)
             self.pipe.to(self.device)
         else:
-            self.pipe = UniDiffuserPipeline.from_pretrained(
-                'thu-ml/unidiffuser-v1')
+            self.pipe = UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1")
 
     def run(
         self,
@@ -27,61 +24,55 @@ class Model:
         guidance_scale: float = 8.0,
     ) -> tuple[PIL.Image.Image | None, str]:
         generator = torch.Generator(device=self.device).manual_seed(seed)
-        if mode == 't2i':
+        if mode == "t2i":
             self.pipe.set_text_to_image_mode()
-            sample = self.pipe(prompt=prompt,
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
-            return sample.images[0], ''
-        elif mode == 'i2t':
+            sample = self.pipe(
+                prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
+            )
+            return sample.images[0], ""
+        elif mode == "i2t":
             self.pipe.set_image_to_text_mode()
-            sample = self.pipe(image=image,
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(
+                image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
+            )
             return None, sample.text[0]
-        elif mode == 'joint':
+        elif mode == "joint":
             self.pipe.set_joint_mode()
-            sample = self.pipe(num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
             return sample.images[0], sample.text[0]
-        elif mode == 'i':
+        elif mode == "i":
             self.pipe.set_image_mode()
-            sample = self.pipe(num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
-            return sample.images[0], ''
-        elif mode == 't':
+            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
+            return sample.images[0], ""
+        elif mode == "t":
             self.pipe.set_text_mode()
-            sample = self.pipe(num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator)
             return None, sample.text[0]
-        elif mode == 'i2t2i':
+        elif mode == "i2t2i":
             self.pipe.set_image_to_text_mode()
-            sample = self.pipe(image=image,
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(
+                image=image, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
+            )
             self.pipe.set_text_to_image_mode()
-            sample = self.pipe(prompt=sample.text[0],
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
-            return sample.images[0], ''
-        elif mode == 't2i2t':
+            sample = self.pipe(
+                prompt=sample.text[0],
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            )
+            return sample.images[0], ""
+        elif mode == "t2i2t":
             self.pipe.set_text_to_image_mode()
-            sample = self.pipe(prompt=prompt,
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(
+                prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=generator
+            )
             self.pipe.set_image_to_text_mode()
-            sample = self.pipe(image=sample.images[0],
-                               num_inference_steps=num_steps,
-                               guidance_scale=guidance_scale,
-                               generator=generator)
+            sample = self.pipe(
+                image=sample.images[0],
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            )
             return None, sample.text[0]
         else:
             raise ValueError