Initial commit

Author: stefan-baumann

Diff for: .gitignore

@@ -0,0 +1,163 @@
+# 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/
+
+# Hydra output dir
+outputs/

Diff for: LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 CompVis - Computer Vision and Learning LMU Munich
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Diff for: README.md

@@ -0,0 +1,67 @@
+# Fine-Grained Subject-Specific Attribute Expression Control in T2I Models
+
+[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://compvis.github.io/attribute-control/)
+[![Paper](https://img.shields.io/badge/arXiv-PDF-b31b1b)](https://arxiv.org/abs/TODO)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/CompVis/attribute-control/blob/main/notebooks/inference_sdxl_colab.ipynb)
+
+This repository contains an implementation of the paper "Continuous, Subject-Specific Attribute Control in T2I Models by Identifying Semantic Directions".
+
+[Stefan Andreas Baumann](https://stefan-baumann.eu/), [Felix Krause](https://www.linkedin.com/in/felixmkrause/), [Michael Neumayr](https://www.linkedin.com/in/michaelneumayr/), [Nick Stracke](https://de.linkedin.com/in/nick-stracke), [Vincent Tao Hu](https://taohu.me/), [Björn Ommer](https://ommer-lab.com/people/ommer/)
+
+We present a simple, straight-forward method for enabling fine-grained control over attribute expression in T2I (diffusion) models in a subject-specific manner.
+We identify meaningful directions in the tokenwise prompt embedding space that enable modulating single attributes for specific subjects without adapting the T2I model.
+
+![teaser](./docs/static/images/teaser.png)
+
+## 🚀 Usage
+### Setup
+Just clone the repo and install the requirements via `pip install -r requirements.txt`, then you're ready to go. For usage, see the examples below, everything else that's needed (model checkpoints) will be downloaded automatically.
+
+### Inference
+For inference, just start with one of the notebook at `notebooks` for a minimal example.
+
+We provide a range of learned deltas for SDXL at `pretrained_deltas`. These can also be used for models such as SD 1.5 or LDM3D by just loading them as usual .
+
+### Creating new Attribute Deltas
+When creating deltas for new attributes, start by creating a config for them akin to `configs/prompts/people/age.yaml`. There are multiple entries of base prompts that correspond to the attribute in a neutral, "negative", and "positive" direction. Please make sure to use the same noun for all the prompts per entry and specify it as the `pattern_target`.
+You can also specify a list of prefixes that contain various other words that will be added before the main prompt to help obtain more robust deltas. The syntax used finds all sets of words enclosed in braces (e.g., `{young,old}`) and then generates all combinations of words in the braces.
+
+#### Learning-based Method
+The best method to obtain deltas is the learning-based method, although it takes substantially longer than the naive method (see below)
+
+To obtain a delta with the naive method, use:
+```shell
+python learn_delta.py device=cuda:0 model=sdxl prompts=people/age
+```
+This will save the delta at `outputs/learn_delta/people/age/runs/<date>/<time>/checkpoints/delta.pt`, which you can then directly use as shown in the example notebooks.
+
+#### Naive CLIP Difference Method
+The simplest method to obtain deltas is the naive CLIP difference-based method. With it, you can obtain a delta in a few seconds on a decent GPU. It is substantially worse than the proper learned method though.
+
+To obtain a delta with the naive method, use (same arguments as for the learning-based method):
+```shell
+python learn_delta_naive_clip.py device=cuda:0 model=sdxl prompts=people/age
+```
+This will save the delta at `outputs/learn_delta_naive_clip/people/age/runs/<date>/<time>/checkpoints/delta.pt`, which you can then directly use as shown in the example notebooks.
+
+## Future Plans
+This repository contains a clean re-implementation of the code used to create our paper. Therefore, it is still missing some non-essential features. We are planning to add these in the near future.
+
+We also hope to add support for inference & delta learning with more models in the future.
+
+If you have any suggestions as to what you'd like to see, let us know in the issues!
+
+We also welcome external contributions! Additionally, if you build something cool with this, let us know so that we can add a link here.
+
+## 🎓 Citation
+
+If you use this codebase or otherwise found our work valuable, please cite our paper:
+
+```bibtex
+@misc{baumann2024attributecontrol,
+      title={Continuous, Subject-Specific Attribute Control in T2I Models by Identifying Semantic Directions}, 
+      author={Baumann, Stefan Andreas and Krause, Felix and Neumayr, Michael and Stracke, Nick and Hu, Vincent Tao and Ommer, Bj\"orn},
+      year={2024},
+      archivePrefix={arXiv},
+}
+```

Diff for: attribute_control/__init__.py

@@ -0,0 +1,2 @@
+from . import utils
+from .base import PromptEmbedding, EmbeddingDelta

Diff for: attribute_control/base.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Tuple, Dict, List
+import torch
+from torch import nn
+from jaxtyping import Float
+
+
+@dataclass
+class PromptEmbedding:
+    prompt: str
+    tokenwise_embeddings: Dict[str, Float[torch.Tensor, 'n d']]
+    tokenwise_embedding_spans: Dict[str, List[Tuple[int, int]]]
+    pooled_embeddings: Dict[str, Float[torch.Tensor, 'd']] = field(default_factory=dict)
+
+    @staticmethod
+    def merge(*embs: PromptEmbedding) -> PromptEmbedding:
+        emb_joined = None
+        for emb in embs:
+            if emb_joined is None:
+                emb_joined = PromptEmbedding(
+                    prompt=emb.prompt,
+                    tokenwise_embeddings={ k: v for k, v in emb.tokenwise_embeddings.items() },
+                    tokenwise_embedding_spans={ k: v for k, v in emb.tokenwise_embedding_spans.items() },
+                    pooled_embeddings={ k: v for k, v in emb.pooled_embeddings.items() },
+                )
+            else:
+                assert emb.prompt == emb_joined.prompt
+                emb_joined.tokenwise_embeddings = emb_joined.tokenwise_embeddings | emb.tokenwise_embeddings
+                emb_joined.tokenwise_embedding_spans = emb_joined.tokenwise_embedding_spans | emb.tokenwise_embedding_spans
+                emb_joined.pooled_embeddings = emb_joined.pooled_embeddings | emb.pooled_embeddings
+        return emb_joined
+
+    def get_tokenwise_mask(self, characterwise_mask: List[bool]) -> Dict[str, List[bool]]:
+        tokenwise_masks = { }
+        for k, t_embs, t_spans in ((k, self.tokenwise_embeddings[k], self.tokenwise_embedding_spans[k]) for k in self.tokenwise_embeddings):
+            token_mask = [False] * len(t_embs)
+            for i_t, (t_span_start, t_span_end) in enumerate(t_spans):
+                if t_span_start != t_span_end: # Don't apply to SOS/EOS tokens
+                    m = characterwise_mask[t_span_start:t_span_end]
+                    assert all(m) or not any(m), 'Inconsistent mask'
+                    if all(m):
+                        token_mask[i_t] = True
+            tokenwise_masks[k] = token_mask
+        return tokenwise_masks
+
+
+class EmbeddingDelta(nn.Module):
+    def __init__(self, dims: Dict[str, int]) -> None:
+        super().__init__()
+        self.tokenwise_delta = nn.ParameterDict({ k: nn.Parameter(torch.zeros(d), requires_grad=True) for k, d in dims.items() })
+    
+    def apply(self, emb: PromptEmbedding, characterwise_mask: List[bool], alpha: float = 1.) -> PromptEmbedding:
+        tokenwise_embeddings = { }
+        matching_keys = [k for k in self.tokenwise_delta if k in emb.tokenwise_embeddings]
+        assert len(matching_keys) > 0, f'Cannot apply delta if no embeddings match. Embeddings present in delta: {[k for k in self.tokenwise_delta]}. Embeddings present in embedding: {[k for k in emb.tokenwise_embeddings]}.'
+        for k, t_embs, t_spans in ((k, emb.tokenwise_embeddings[k], emb.tokenwise_embedding_spans[k]) for k in matching_keys):
+            token_mask = [0] * len(t_embs)
+            for i_t, (t_span_start, t_span_end) in enumerate(t_spans):
+                if t_span_start != t_span_end: # Don't apply to SOS/EOS tokens
+                    m = characterwise_mask[t_span_start:t_span_end]
+                    assert all(m) or not any(m), 'Inconsistent mask'
+                    if all(m):
+                        token_mask[i_t] = 1
+            assert sum(token_mask) >= 1, f'No tokens in prompt selected for delta application via characterwise mask.'
+            tokenwise_embeddings[k] = t_embs + alpha * torch.tensor(token_mask, dtype=t_embs.dtype, device=t_embs.device).unsqueeze(-1) * self.tokenwise_delta[k].unsqueeze(0).to(t_embs.dtype)
+        return PromptEmbedding(
+            prompt=emb.prompt,
+            tokenwise_embeddings=tokenwise_embeddings,
+            tokenwise_embedding_spans=emb.tokenwise_embedding_spans,
+            pooled_embeddings=emb.pooled_embeddings
+        )

Diff for: attribute_control/model/__init__.py

@@ -0,0 +1 @@
+from .model import ModelBase, SD15, SDXL, StableCascade