Add files

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "plug-and-play"]
+	path = plug-and-play
+	url = https://github.com/MichalGeyer/plug-and-play

.pre-commit-config.yaml

@@ -0,0 +1,37 @@
+exclude: patch
+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']

.style.yapf

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

README.md

@@ -5,6 +5,7 @@ colorFrom: blue
 colorTo: pink
 sdk: gradio
 sdk_version: 3.18.0
+python_version: 3.10.9
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import pathlib
+import shlex
+import subprocess
+
+import gradio as gr
+
+from app_generated_image import create_prompt_demo
+from app_real_image import create_real_image_demo
+
+DESCRIPTION = '''# Plug-and-Play diffusion features
+
+This is an unofficial demo for [https://github.com/MichalGeyer/plug-and-play](https://github.com/MichalGeyer/plug-and-play).
+'''
+
+weight_dir = pathlib.Path('plug-and-play/models/ldm/stable-diffusion-v1')
+if not weight_dir.exists():
+    subprocess.run(
+        shlex.split('mkdir -p plug-and-play/models/ldm/stable-diffusion-v1/'))
+    subprocess.run(
+        shlex.split(
+            'wget https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt -O plug-and-play/models/ldm/stable-diffusion-v1/model.ckpt'
+        ))
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Tabs():
+        with gr.TabItem('Use real image as input'):
+            create_real_image_demo()
+        with gr.TabItem('Use prompt as input'):
+            create_prompt_demo()
+
+demo.queue().launch()

app_generated_image.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import pathlib
+import shlex
+import subprocess
+import tempfile
+
+import gradio as gr
+from omegaconf import OmegaConf
+
+from utils import get_timestamp
+
+
+def gen_feature_extraction_config(
+    exp_name: str,
+    prompt: str,
+    seed: int,
+    guidance_scale: float,
+    ddim_steps: int,
+) -> str:
+    config = OmegaConf.load(
+        'plug-and-play/configs/pnp/feature-extraction-generated.yaml')
+    config.config.experiment_name = exp_name
+    config.config.prompt = prompt
+    config.config.seed = seed
+    config.config.scale = guidance_scale
+    config.config.ddim_steps = ddim_steps
+    temp_file = tempfile.NamedTemporaryFile(suffix='.yaml', delete=False)
+    with open(temp_file.name, 'w') as f:
+        f.write(OmegaConf.to_yaml(config))
+    return temp_file.name
+
+
+def run_feature_extraction_command(
+    prompt: str,
+    seed: int,
+    guidance_scale: float,
+    ddim_steps: int,
+) -> tuple[str, str]:
+    exp_name = get_timestamp()
+    config_path = gen_feature_extraction_config(
+        exp_name,
+        prompt,
+        seed,
+        guidance_scale,
+        ddim_steps,
+    )
+    subprocess.run(shlex.split(
+        f'python run_features_extraction.py --config {config_path}'),
+                   cwd='plug-and-play')
+    return f'plug-and-play/experiments/{exp_name}/samples/0.png', exp_name
+
+
+def gen_pnp_config(
+    exp_name: str,
+    prompt: str,
+    guidance_scale: float,
+    ddim_steps: int,
+    feature_injection_threshold: int,
+    negative_prompt: str,
+    negative_prompt_alpha: float,
+    negative_prompt_schedule: str,
+) -> str:
+    config = OmegaConf.load('plug-and-play/configs/pnp/pnp-generated.yaml')
+    config.source_experiment_name = exp_name
+    config.prompts = [prompt]
+    config.scale = guidance_scale
+    config.num_ddim_sampling_steps = ddim_steps
+    config.feature_injection_threshold = feature_injection_threshold
+    config.negative_prompt = negative_prompt
+    config.negative_prompt_alpha = negative_prompt_alpha
+    config.negative_prompt_schedule = negative_prompt_schedule
+    temp_file = tempfile.NamedTemporaryFile(suffix='.yaml', delete=False)
+    with open(temp_file.name, 'w') as f:
+        f.write(OmegaConf.to_yaml(config))
+    return temp_file.name
+
+
+def run_pnp_command(
+    exp_name: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float,
+    ddim_steps: int,
+    feature_injection_threshold: int,
+    negative_prompt_alpha: float,
+    negative_prompt_schedule: str,
+) -> str:
+    config_path = gen_pnp_config(
+        exp_name,
+        prompt,
+        guidance_scale,
+        ddim_steps,
+        feature_injection_threshold,
+        negative_prompt,
+        negative_prompt_alpha,
+        negative_prompt_schedule,
+    )
+    subprocess.run(shlex.split(f'python run_pnp.py --config {config_path}'),
+                   cwd='plug-and-play')
+
+    out_dir = pathlib.Path(
+        f'plug-and-play/experiments/{exp_name}/translations/{guidance_scale}_{prompt.replace(" ", "_")}'
+    )
+    out_label = f'INJECTION_T_{feature_injection_threshold}_STEPS_{ddim_steps}_NP-ALPHA_{negative_prompt_alpha}_SCHEDULE_{negative_prompt_schedule}_NP_{negative_prompt.replace(" ", "_")}'
+    out_path = out_dir / f'{out_label}_sample_0.png'
+    return out_path.as_posix()
+
+
+def create_prompt_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        with gr.Box():
+            gr.Markdown('Step 1')
+            with gr.Row():
+                with gr.Column():
+                    source_prompt = gr.Text(label='Source prompt')
+                    seed = gr.Slider(label='Seed',
+                                     minimum=0,
+                                     maximum=100000,
+                                     step=1,
+                                     value=0)
+                    with gr.Accordion(label='Advanced settings', open=False):
+                        source_guidance_scale = gr.Slider(
+                            label='Guidance scale',
+                            minimum=0,
+                            maximum=50,
+                            step=0.1,
+                            value=5)
+                        source_ddim_steps = gr.Slider(label='DDIM steps',
+                                                      minimum=1,
+                                                      maximum=100,
+                                                      step=1,
+                                                      value=50)
+                    extract_feature_button = gr.Button(
+                        'Generate and extract features')
+                with gr.Column():
+                    generated_image = gr.Image(label='Generated image',
+                                               type='filepath')
+                    exp_name = gr.Variable()
+        with gr.Box():
+            gr.Markdown('Step 2')
+            with gr.Row():
+                with gr.Column():
+                    translation_prompt = gr.Text(
+                        label='Prompt for translation')
+                    negative_prompt = gr.Text(label='Negative prompt')
+                    with gr.Accordion(label='Advanced settings', open=False):
+                        guidance_scale = gr.Slider(label='Guidance scale',
+                                                   minimum=0,
+                                                   maximum=50,
+                                                   step=0.1,
+                                                   value=7.5)
+                        ddim_steps = gr.Slider(
+                            label='Number of inference steps',
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=50)
+                        feature_injection_threshold = gr.Slider(
+                            label='Feature injection threshold',
+                            minimum=0,
+                            maximum=100,
+                            step=1,
+                            value=40)
+                        negative_prompt_alpha = gr.Slider(
+                            label='Negative prompt alpha',
+                            minimum=0,
+                            maximum=1,
+                            step=0.01,
+                            value=0.75)
+                        negative_prompt_schedule = gr.Dropdown(
+                            label='Negative prompt schedule',
+                            choices=['linear', 'constant', 'exp'],
+                            value='linear')
+                    generate_button = gr.Button('Generate')
+                with gr.Column():
+                    result = gr.Image(label='Result', type='filepath')
+        with gr.Row():
+            gr.Examples(examples=[
+                ['horse in mud', 50, 'a photo of a zebra in the snow'],
+                ['horse in mud', 50, 'a photo of a husky in the grass'],
+            ],
+                        inputs=[
+                            source_prompt,
+                            seed,
+                            translation_prompt,
+                        ])
+
+        extract_feature_button.click(
+            fn=run_feature_extraction_command,
+            inputs=[
+                source_prompt,
+                seed,
+                source_guidance_scale,
+                source_ddim_steps,
+            ],
+            outputs=[
+                generated_image,
+                exp_name,
+            ],
+        )
+        generate_button.click(
+            fn=run_pnp_command,
+            inputs=[
+                exp_name,
+                translation_prompt,
+                negative_prompt,
+                guidance_scale,
+                ddim_steps,
+                feature_injection_threshold,
+                negative_prompt_alpha,
+                negative_prompt_schedule,
+            ],
+            outputs=result,
+        )
+    return demo
+
+
+if __name__ == '__main__':
+    demo = create_prompt_demo()
+    demo.queue().launch()

app_real_image.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import pathlib
+import shlex
+import subprocess
+import tempfile
+
+import gradio as gr
+from omegaconf import OmegaConf
+
+from utils import get_timestamp
+
+
+def gen_feature_extraction_config(exp_name: str, init_image_path: str) -> str:
+    config = OmegaConf.load(
+        'plug-and-play/configs/pnp/feature-extraction-real.yaml')
+    config.config.experiment_name = exp_name
+    config.config.init_img = init_image_path
+    temp_file = tempfile.NamedTemporaryFile(suffix='.yaml', delete=False)
+    with open(temp_file.name, 'w') as f:
+        f.write(OmegaConf.to_yaml(config))
+    return temp_file.name
+
+
+def run_feature_extraction_command(init_image_path: str) -> tuple[str, str]:
+    exp_name = get_timestamp()
+    config_path = gen_feature_extraction_config(exp_name, init_image_path)
+    subprocess.run(shlex.split(
+        f'python run_features_extraction.py --config {config_path}'),
+                   cwd='plug-and-play')
+    return f'plug-and-play/experiments/{exp_name}/samples/0.png', exp_name
+
+
+def gen_pnp_config(
+    exp_name: str,
+    prompt: str,
+    guidance_scale: float,
+    ddim_steps: int,
+    feature_injection_threshold: int,
+    negative_prompt: str,
+    negative_prompt_alpha: float,
+    negative_prompt_schedule: str,
+) -> str:
+    config = OmegaConf.load('plug-and-play/configs/pnp/pnp-real.yaml')
+    config.source_experiment_name = exp_name
+    config.prompts = [prompt]
+    config.scale = guidance_scale
+    config.num_ddim_sampling_steps = ddim_steps
+    config.feature_injection_threshold = feature_injection_threshold
+    config.negative_prompt = negative_prompt
+    config.negative_prompt_alpha = negative_prompt_alpha
+    config.negative_prompt_schedule = negative_prompt_schedule
+    temp_file = tempfile.NamedTemporaryFile(suffix='.yaml', delete=False)
+    with open(temp_file.name, 'w') as f:
+        f.write(OmegaConf.to_yaml(config))
+    return temp_file.name
+
+
+def run_pnp_command(
+    exp_name: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float,
+    ddim_steps: int,
+    feature_injection_threshold: int,
+    negative_prompt_alpha: float,
+    negative_prompt_schedule: str,
+) -> str:
+    config_path = gen_pnp_config(
+        exp_name,
+        prompt,
+        guidance_scale,
+        ddim_steps,
+        feature_injection_threshold,
+        negative_prompt,
+        negative_prompt_alpha,
+        negative_prompt_schedule,
+    )
+    subprocess.run(shlex.split(f'python run_pnp.py --config {config_path}'),
+                   cwd='plug-and-play')
+
+    out_dir = pathlib.Path(
+        f'plug-and-play/experiments/{exp_name}/translations/{guidance_scale}_{prompt.replace(" ", "_")}'
+    )
+    out_label = f'INJECTION_T_{feature_injection_threshold}_STEPS_{ddim_steps}_NP-ALPHA_{negative_prompt_alpha}_SCHEDULE_{negative_prompt_schedule}_NP_{negative_prompt.replace(" ", "_")}'
+    out_path = out_dir / f'{out_label}_sample_0.png'
+    return out_path.as_posix()
+
+
+def create_real_image_demo():
+    with gr.Blocks() as demo:
+        with gr.Box():
+            gr.Markdown('Step 1')
+            with gr.Row():
+                with gr.Column():
+                    image = gr.Image(label='Input image', type='filepath')
+                    extract_feature_button = gr.Button(
+                        'Reconstruct and extract features')
+                with gr.Column():
+                    reconstructed_image = gr.Image(label='Reconstructed image',
+                                                   type='filepath')
+                    exp_name = gr.Variable()
+        with gr.Box():
+            gr.Markdown('Step 2')
+            with gr.Row():
+                with gr.Column():
+                    translation_prompt = gr.Text(
+                        label='Prompt for translation')
+                    negative_prompt = gr.Text(label='Negative prompt')
+                    with gr.Accordion(label='Advanced settings', open=False):
+                        guidance_scale = gr.Slider(label='Guidance scale',
+                                                   minimum=0,
+                                                   maximum=50,
+                                                   step=0.1,
+                                                   value=10)
+                        ddim_steps = gr.Slider(
+                            label='Number of inference steps',
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=50)
+                        feature_injection_threshold = gr.Slider(
+                            label='Feature injection threshold',
+                            minimum=0,
+                            maximum=100,
+                            step=1,
+                            value=40)
+                        negative_prompt_alpha = gr.Slider(
+                            label='Negative prompt alpha',
+                            minimum=0,
+                            maximum=1,
+                            step=0.01,
+                            value=1)
+                        negative_prompt_scheduler = gr.Dropdown(
+                            label='Negative prompt schedule',
+                            choices=['linear', 'constant', 'exp'],
+                            value='linear')
+                    generate_button = gr.Button('Generate')
+                with gr.Column():
+                    result = gr.Image(label='Result', type='filepath')
+
+        with gr.Row():
+            gr.Examples(examples=[
+                [
+                    'plug-and-play/data/horse.png',
+                    'a photo of a robot horse',
+                    'a photo of a white horse',
+                ],
+                [
+                    'plug-and-play/data/horse.png',
+                    'a photo of a bronze horse in a museum',
+                    'a photo of a white horse',
+                ],
+            ],
+                        inputs=[
+                            image,
+                            translation_prompt,
+                            negative_prompt,
+                        ])
+
+        extract_feature_button.click(
+            fn=run_feature_extraction_command,
+            inputs=image,
+            outputs=[
+                reconstructed_image,
+                exp_name,
+            ],
+        )
+        generate_button.click(
+            fn=run_pnp_command,
+            inputs=[
+                exp_name,
+                translation_prompt,
+                negative_prompt,
+                guidance_scale,
+                ddim_steps,
+                feature_injection_threshold,
+                negative_prompt_alpha,
+                negative_prompt_scheduler,
+            ],
+            outputs=result,
+        )
+
+    return demo
+
+
+if __name__ == '__main__':
+    demo = create_real_image_demo()
+    demo.queue().launch()

packages.txt

@@ -0,0 +1,2 @@
+ffmpeg
+wget

plug-and-play

@@ -0,0 +1 @@
+Subproject commit ce9e4e23c24b4241cb9ecae9ea0b090c70871870

requirements.txt

@@ -0,0 +1,19 @@
+albumentations==1.3.0
+-e git+https://github.com/openai/CLIP.git@main#egg=clip
+diffusers==0.12.1
+einops==0.6.0
+gradio==3.18.0
+huggingface-hub==0.12.0
+imageio==2.25.0
+imageio-ffmpeg==0.4.8
+kornia==0.6.9
+omegaconf==2.3.0
+opencv-python-headless==4.7.0.68
+Pillow==9.4.0
+python-slugify==7.0.0
+pytorch-lightning==1.9.0
+scikit-learn==1.2.0
+-e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+torch==1.13.1
+torchvision==0.14.1
+transformers==4.26.0

style.css

@@ -0,0 +1,3 @@
+h1 {
+  text-align: center;
+}

utils.py

@@ -0,0 +1,5 @@
+import datetime
+
+
+def get_timestamp() -> str:
+    return datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S-%f')