Initial commit

.gitignore

@@ -0,0 +1,180 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+.vscode/
+*.py[cod]
+*$py.class
+*.csv
+*.png
+*.jpg
+*.parquet
+# Folders
+models/
+pickles/
+train_log/
+test_log/
+analysis/
+classification_models/
+nlp_models/
+# C extensions
+*.so
+.idea
+.DS_Store
+tokens.json
+testing.py
+# Distribution / packaging
+Git
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+tokens.json
+*.pkl
+test.py
+
+# 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/
+INSTALL.md
+
+# 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/

README.md

@@ -0,0 +1,63 @@
+# TryOffAnyone
+
+The official repository of the paper: _"TryOffAnyone: Tiled Cloth Generation from a Dressed Person"_.
+
+[![arXiv][logo-paper]][paper-arxiv]
+[![Generic badge][logo-hf_models]][hf_models]
+
+
+### Install
+Clone the repository, install the required packages:
+```
+git clone https://github.com/ixarchakos/try-off-anyone.git
+cd try-off-anyone
+python3 -m pip install -r requirements.txt
+```
+
+Download the model from:
+- https://huggingface.co/ixarchakos/tryOffAnyone
+
+in `"try-off-anyone/ckpt/"`
+
+The provided model only works for upper clothes similar to the ones in VITON-HD
+
+### Inference
+Execute TryOffAnyone on any image url
+```
+python3 main.py --inference --url="https://cdn11.bigcommerce.com/s-405b0/images/stencil/590x590/products/97/20409/8000-gildan-tee-t-shirt.ca-model__66081.1724276210.jpg"
+```
+
+the output is saved in `"try-off-anyone/data/"`
+
+### Evaluation on VITON-HD
+
+#### Dataset
+- Download the original VITON-HD dataset 
+
+and extract it to `"try-off-anyone/data/"`
+
+- Download the cloth image masks from:
+https://huggingface.co/ixarchakos/tryOffAnyone/blob/main/masks.zip
+
+and save them in `"try-off-anyone/data/zalando-hd-resized/test/"`
+
+#### Evaluation
+
+Evaluate TryOffAnyone on VITON-HD:
+```bash
+python3 main.py --test
+```
+
+
+### Citation
+If you find this repository useful in your research, please consider giving a star ⭐ and a citation:
+```
+
+```
+
+[logo-hf_models]: https://img.shields.io/badge/🤗-Models-blue.svg?style=plastic
+[logo-hf_spaces]: https://img.shields.io/badge/🤗-Demo-blue.svg?style=plastic
+[logo-paper]: https://img.shields.io/badge/arXiv-Paper-b31b1b.svg?style=plastic
+[hf_datasets]: https://huggingface.co/datasets/rizavelioglu/...
+[hf_models]: https://huggingface.co/ixarchakos/tryOffAnyone
+[paper-arxiv]: https://arxiv.org/

config.py

@@ -0,0 +1,7 @@
+import torch
+
+TEST_DATA_PATH = 'data/zalando-hd-resized/test/'
+device = 'cuda'
+concat_d = -2
+dtype = torch.bfloat16
+base_ckpt = "runwayml/stable-diffusion-inpainting"

main.py

@@ -0,0 +1,22 @@
+from src.test_vton import test_vton
+from src.inference import test_image
+import argparse
+
+
+def terminal_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument("--test", action="store_true")
+    parser.add_argument("--inference", action="store_true")
+    return parser.parse_known_args()[0]
+
+
+def main():
+    args = terminal_args()
+    if args.test:
+        test_vton()
+    elif args.inference:
+        test_image()
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,16 @@
+torch
+safetensors
+tqdm
+diffusers
+pillow
+numpy
+requests
+datasets
+transformers
+opencv-python
+torchvision
+psutil
+scipy
+setuptools
+torchmetrics
+torch-fidelity

src/inference.py

@@ -0,0 +1,66 @@
+from config import device
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from diffusers.image_processor import VaeImageProcessor
+from src.model.pipeline import TryOffAnyone
+from src.preprocessing import background_removal, background_whitening, mask_generation
+from PIL import Image
+from io import BytesIO
+import argparse
+import torch
+import os
+import requests
+
+
+def inference_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seed", type=int, default=36)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--scale", type=float, default=2.5)
+    parser.add_argument("--width", type=int, default=384)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument('--gpu_id', type=int, default=0)
+    parser.add_argument('--url', type=str, default="https://cdn11.bigcommerce.com/s-405b0/images/stencil/590x590/products/97/20409/8000-gildan-tee-t-shirt.ca-model__66081.1724276210.jpg")
+    return parser.parse_known_args()[0]
+
+
+def get_image_file(image_url):
+    try:
+        image_response = requests.get(image_url, timeout=20)
+    except Exception:
+        raise "Provide valid url!"
+    image_response.raise_for_status()
+    return BytesIO(image_response.content)
+
+
+def generate_laydown(pipeline, cloth_image, mask, args):
+    result = pipeline(
+        cloth_image, mask, inference_steps=args.steps, scale=args.scale, height=args.height,
+        width=args.width, generator=torch.Generator(device=device).manual_seed(args.seed)
+    )
+    return result[0]
+
+
+def test_image():
+    args = inference_args()
+    pipeline = TryOffAnyone()
+    processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer_b3_clothes")
+    model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer_b3_clothes")
+    model.to(device)
+    mask_processor = VaeImageProcessor(
+        vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True
+    )
+    vae_processor = VaeImageProcessor(vae_scale_factor=8)
+
+    image = Image.open(get_image_file(args.url))
+    image = image.convert("RGB").resize((args.width, args.height))
+
+    mask_image = mask_generation(image, processor, model, "Tops")
+
+    mask = mask_processor.preprocess(mask_image, args.height, args.width)[0]
+    image = vae_processor.preprocess(image, args.height, args.width)[0]
+
+    laydown_image = generate_laydown(pipeline, image, mask, args)
+    laydown_image = background_whitening(
+        background_removal(laydown_image), args.width, args.height
+    )
+    laydown_image.save(os.path.join('data', f"{args.url.split('/')[-1][:-4]}.png"))

src/model/attention.py

@@ -0,0 +1,11 @@
+import torch
+
+
+class Skip(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __call__(
+            self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None
+    ):
+        return hidden_states