Duplicate from nupurkmr9/custom-diffusion

Co-authored-by: Nupur Kumari <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+method.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,164 @@
+training_data/
+results/
+
+
+# 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 "custom-diffusion"]
+	path = custom-diffusion
+	url = https://github.com/adobe-research/custom-diffusion

.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+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.10.1
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+- 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

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 hysts
+
+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.

README.md

@@ -0,0 +1,13 @@
+---
+title: Custom-Diffusion + SD Training
+emoji: 🏢
+colorFrom: red
+colorTo: purple
+sdk: gradio
+sdk_version: 3.12.0
+app_file: app.py
+pinned: false
+duplicated_from: nupurkmr9/custom-diffusion
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python
+"""Demo app for https://github.com/adobe-research/custom-diffusion.
+
+The code in this repo is partly adapted from the following repository:
+https://huggingface.co/spaces/hysts/LoRA-SD-training
+"""
+
+from __future__ import annotations
+import sys
+import os
+import pathlib
+
+import gradio as gr
+import torch
+
+from inference import InferencePipeline
+from trainer import Trainer
+from uploader import upload
+
+TITLE = '# Custom Diffusion + StableDiffusion Training UI'
+DESCRIPTION = '''This is a demo for [https://github.com/adobe-research/custom-diffusion](https://github.com/adobe-research/custom-diffusion).
+It is recommended to upgrade to GPU in Settings after duplicating this space to use it.
+<a href="https://huggingface.co/spaces/nupurkmr9/custom-diffusion?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+'''
+DETAILDESCRIPTION='''
+Custom Diffusion allows you to fine-tune text-to-image diffusion models, such as Stable Diffusion, given a few images of a new concept (~4-20). 
+We fine-tune only a subset of model parameters, namely key and value projection matrices, in the cross-attention layers and the modifier token used to represent the object. 
+This also reduces the extra storage for each additional concept to 75MB.
+Our method further allows you to use a combination of concepts. Demo for multiple concepts will be added soon.
+<center>
+<img src="https://huggingface.co/spaces/nupurkmr9/custom-diffusion/resolve/main/method.jpg" width="600" align="center" >
+</center>
+'''
+
+ORIGINAL_SPACE_ID = 'nupurkmr9/custom-diffusion'
+SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
+SHARED_UI_WARNING = f'''# Attention - This Space doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
+
+<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center>
+'''
+if os.getenv('SYSTEM') == 'spaces' and SPACE_ID != ORIGINAL_SPACE_ID:
+    SETTINGS = f'<a href="https://huggingface.co/spaces/{SPACE_ID}/settings">Settings</a>'
+
+else:
+    SETTINGS = 'Settings'
+CUDA_NOT_AVAILABLE_WARNING = f'''# Attention - Running on CPU.
+<center>
+You can assign a GPU in the {SETTINGS} tab if you are running this on HF Spaces.
+"T4 small" is sufficient to run this demo.
+</center>
+'''
+
+os.system("git clone https://github.com/adobe-research/custom-diffusion")
+sys.path.append("custom-diffusion")
+
+def show_warning(warning_text: str) -> gr.Blocks:
+    with gr.Blocks() as demo:
+        with gr.Box():
+            gr.Markdown(warning_text)
+    return demo
+
+
+def update_output_files() -> dict:
+    paths = sorted(pathlib.Path('results').glob('*.pt'))
+    paths = [path.as_posix() for path in paths]  # type: ignore
+    return gr.update(value=paths or None)
+
+
+def create_training_demo(trainer: Trainer,
+                         pipe: InferencePipeline) -> gr.Blocks:
+    with gr.Blocks() as demo:
+        base_model = gr.Dropdown(
+            choices=['stabilityai/stable-diffusion-2-1-base', 'CompVis/stable-diffusion-v1-4'],
+            value='CompVis/stable-diffusion-v1-4',
+            label='Base Model',
+            visible=True)
+        resolution = gr.Dropdown(choices=['512', '768'],
+                                 value='512',
+                                 label='Resolution',
+                                 visible=True)
+
+        with gr.Row():
+            with gr.Box():
+                gr.Markdown('Training Data')
+                concept_images = gr.Files(label='Images for your concept')
+                with gr.Row():
+                    class_prompt = gr.Textbox(label='Class Prompt',
+                                            max_lines=1, placeholder='Example: "cat"')
+                    with gr.Column():
+                        modifier_token = gr.Checkbox(label='modifier token',
+                                                    value=True)
+                        train_text_encoder = gr.Checkbox(label='Train Text Encoder',
+                                                 value=False)
+                concept_prompt = gr.Textbox(label='Concept Prompt',
+                                                max_lines=1, placeholder='Example: "photo of a \<new1\> cat"')
+                gr.Markdown('''
+                    - We use "\<new1\>" modifier token in front of the concept, e.g., "\<new1\> cat". By default modifier_token is enabled.
+                    - If "Train Text Encoder", disable "modifier token" and use any unique text to describe the concept e.g. "ktn cat". 
+                    - For a new concept an e.g. concept prompt is "photo of a \<new1\> cat" and "cat" for class prompt.
+                    - For a style concept, use "painting in the style of \<new1\> art" for concept prompt and "art" for class prompt.
+                    - Class prompt should be the object category.
+                    ''')
+            with gr.Box():
+                gr.Markdown('Training Parameters')
+                num_training_steps = gr.Number(
+                    label='Number of Training Steps', value=1000, precision=0)
+                learning_rate = gr.Number(label='Learning Rate', value=0.00001)
+                batch_size = gr.Number(
+                    label='batch_size', value=1, precision=0)
+                with gr.Row():
+                    use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=True) 
+                    gradient_checkpointing = gr.Checkbox(label='Enable gradient checkpointing', value=False)
+                with gr.Accordion('Other Parameters', open=False):
+                    gradient_accumulation = gr.Number(
+                        label='Number of Gradient Accumulation',
+                        value=1,
+                        precision=0)
+                    gen_images = gr.Checkbox(label='Generated images as regularization',
+                                                 value=False)
+                gr.Markdown('''
+                    - It will take about ~10 minutes to train for 1000 steps and ~21GB on a 3090 GPU. 
+                    - Our results in the paper are trained with batch-size 4 (8 including class regularization samples).
+                    - Enable gradient checkpointing for lower memory requirements (~14GB) at the expense of slower backward pass.
+                    - Note that your trained models will be deleted when the second training is started. You can upload your trained model in the "Upload" tab.
+                    ''')
+
+        run_button = gr.Button('Start Training')
+        with gr.Box():
+            with gr.Row():
+                check_status_button = gr.Button('Check Training Status')
+                with gr.Column():
+                    with gr.Box():
+                        gr.Markdown('Message')
+                        training_status = gr.Markdown()
+                    output_files = gr.Files(label='Trained Weight Files')
+
+        run_button.click(fn=pipe.clear,
+                            inputs=None,
+                            outputs=None,)
+        run_button.click(fn=trainer.run,
+                         inputs=[
+                             base_model,
+                             resolution,
+                             concept_images,
+                             concept_prompt,
+                             class_prompt,
+                             num_training_steps,
+                             learning_rate,
+                             train_text_encoder,
+                             modifier_token,
+                             gradient_accumulation,
+                             batch_size,
+                             use_8bit_adam,
+                             gradient_checkpointing,
+                             gen_images
+                         ],
+                         outputs=[
+                             training_status,
+                             output_files,
+                         ],
+                         queue=False)
+        check_status_button.click(fn=trainer.check_if_running,
+                                  inputs=None,
+                                  outputs=training_status,
+                                  queue=False)
+        check_status_button.click(fn=update_output_files,
+                                  inputs=None,
+                                  outputs=output_files,
+                                  queue=False)
+    return demo
+
+
+def find_weight_files() -> list[str]:
+    curr_dir = pathlib.Path(__file__).parent
+    paths = sorted(curr_dir.rglob('*.bin'))
+    paths = [path for path in paths if '.lfs' not in path.name]
+    return [path.relative_to(curr_dir).as_posix() for path in paths]
+
+
+def reload_custom_diffusion_weight_list() -> dict:
+    return gr.update(choices=find_weight_files())
+
+
+def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                base_model = gr.Dropdown(
+                    choices=['stabilityai/stable-diffusion-2-1-base', 'CompVis/stable-diffusion-v1-4'],
+                    value='CompVis/stable-diffusion-v1-4',
+                    label='Base Model',
+                    visible=True)
+                resolution = gr.Dropdown(choices=[512, 768],
+                                 value=512,
+                                 label='Resolution',
+                                 visible=True)
+                reload_button = gr.Button('Reload Weight List')
+                weight_name = gr.Dropdown(choices=find_weight_files(),
+                                               value='custom-diffusion-models/cat.bin',
+                                               label='Custom Diffusion Weight File')
+                prompt = gr.Textbox(
+                    label='Prompt',
+                    max_lines=1,
+                    placeholder='Example: "\<new1\> cat in outer space"')
+                seed = gr.Slider(label='Seed',
+                                 minimum=0,
+                                 maximum=100000,
+                                 step=1,
+                                 value=42)
+                with gr.Accordion('Other Parameters', open=False):
+                    num_steps = gr.Slider(label='Number of Steps',
+                                          minimum=0,
+                                          maximum=500,
+                                          step=1,
+                                          value=200)
+                    guidance_scale = gr.Slider(label='CFG Scale',
+                                               minimum=0,
+                                               maximum=50,
+                                               step=0.1,
+                                               value=6)
+                    eta = gr.Slider(label='DDIM eta',
+                                               minimum=0,
+                                               maximum=1.,
+                                               step=0.1,
+                                               value=1.)
+                    batch_size = gr.Slider(label='Batch Size',
+                                               minimum=0,
+                                               maximum=10.,
+                                               step=1,
+                                               value=2)
+
+                run_button = gr.Button('Generate')
+
+                gr.Markdown('''
+                - Models with names starting with "custom-diffusion-models/" are the pretrained models provided in the [original repo](https://github.com/adobe-research/custom-diffusion), and the ones with names starting with "results/delta.bin" are your trained models.
+                - After training, you can press "Reload Weight List" button to load your trained model names.
+                - Change default batch-size and steps for faster sampling. 
+                ''')
+            with gr.Column():
+                result = gr.Image(label='Result')
+
+        reload_button.click(fn=reload_custom_diffusion_weight_list,
+                            inputs=None,
+                            outputs=weight_name)
+        prompt.submit(fn=pipe.run,
+                      inputs=[
+                          base_model,
+                          weight_name,
+                          prompt,
+                          seed,
+                          num_steps,
+                          guidance_scale,
+                          eta,
+                          batch_size,
+                          resolution
+                      ],
+                      outputs=result,
+                      queue=False)
+        run_button.click(fn=pipe.run,
+                         inputs=[
+                             base_model,
+                             weight_name,
+                             prompt,
+                             seed,
+                             num_steps,
+                             guidance_scale,
+                             eta,
+                             batch_size,
+                             resolution
+                         ],
+                         outputs=result,
+                         queue=False)
+    return demo
+
+
+def create_upload_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        model_name = gr.Textbox(label='Model Name')
+        hf_token = gr.Textbox(
+            label='Hugging Face Token (with write permission)')
+        upload_button = gr.Button('Upload')
+        with gr.Box():
+            gr.Markdown('Message')
+            result = gr.Markdown()
+        gr.Markdown('''
+            - You can upload your trained model to your private Model repo (i.e. https://huggingface.co/{your_username}/{model_name}).
+            - You can find your Hugging Face token [here](https://huggingface.co/settings/tokens).
+            ''')
+
+    upload_button.click(fn=upload,
+                        inputs=[model_name, hf_token],
+                        outputs=result)
+
+    return demo
+
+
+pipe = InferencePipeline()
+trainer = Trainer()
+
+with gr.Blocks(css='style.css') as demo:
+    if os.getenv('IS_SHARED_UI'):
+        show_warning(SHARED_UI_WARNING)
+    if not torch.cuda.is_available():
+        show_warning(CUDA_NOT_AVAILABLE_WARNING)
+
+    gr.Markdown(TITLE)
+    gr.Markdown(DESCRIPTION)
+    gr.Markdown(DETAILDESCRIPTION)
+
+    with gr.Tabs():
+        with gr.TabItem('Train'):
+            create_training_demo(trainer, pipe)
+        with gr.TabItem('Test'):
+            create_inference_demo(pipe)
+        with gr.TabItem('Upload'):
+            create_upload_demo()
+
+demo.queue(default_enabled=False).launch(share=False)

custom-diffusion-models/barn.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e6ffb7953740286e005fb5ceffc3e985f93b3de97cd46202cf7d66d2171094b
+size 76690626

custom-diffusion-models/cat.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08754e711b9ecaa36785dc64ad0c08317a93d106629c5f42cc5b9a406fe4aefc
+size 76690626

custom-diffusion-models/chair.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e2edb6eaddf540ab9e1a0aa75f3e46ee77c9ee41e8d8e87127777d5dd3ba4b7
+size 76690626

custom-diffusion-models/dog.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dbe8a8843279fa01f2eaa3e9b0b34267e5b1949456f81b5bc17fb2a0d23086fe
+size 76690626

custom-diffusion-models/flower.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a37dcb14359c2baae984a218f62758eb52c842c7557e790063d4cd4daa120e5b
+size 76690626

custom-diffusion-models/moongate.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:865ccdd950d4384af1b4cf45d955db4b26ec3736eb03bccab70fee4f51abb441
+size 76687301

custom-diffusion-models/table.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3771413ef6b319fd1f6df0ec2490febe17524cc3afd67b63042bed85af8cb9c2
+size 76690626

custom-diffusion-models/teddybear.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0656c6eafed89a146f7ab971913c9ee35b4f0a96a4e1aa8eb8ccc28326a8164
+size 76690626

custom-diffusion-models/tortoise_plushy.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38597922778e96c54e3a57379a11a356e62b76be990097922def3f9b764db48d
+size 76690626

custom-diffusion-models/wooden_pot.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:414cbd2bef3e7e65d860c4df17c8f8b8616f5dd8676634e0b228be8ed039f176
+size 76690626

inference.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import gc
+import pathlib
+import sys
+
+import gradio as gr
+import PIL.Image
+import numpy as np
+
+import torch
+from diffusers import StableDiffusionPipeline
+sys.path.insert(0, './custom-diffusion')
+
+
+class InferencePipeline:
+    def __init__(self):
+        self.pipe = None
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.weight_path = None
+
+    def clear(self) -> None:
+        self.weight_path = None
+        del self.pipe
+        self.pipe = None
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    @staticmethod
+    def get_weight_path(name: str) -> pathlib.Path:
+        curr_dir = pathlib.Path(__file__).parent
+        return curr_dir / name
+
+    def load_pipe(self, model_id: str, filename: str) -> None:
+        weight_path = self.get_weight_path(filename)
+        if weight_path == self.weight_path:
+            return
+        self.weight_path = weight_path
+        weight = torch.load(self.weight_path, map_location=self.device)
+
+        if self.device.type == 'cpu':
+            pipe = StableDiffusionPipeline.from_pretrained(model_id)
+        else:
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_id, torch_dtype=torch.float16)
+            pipe = pipe.to(self.device)
+
+        from src import diffuser_training
+        diffuser_training.load_model(pipe.text_encoder, pipe.tokenizer, pipe.unet, weight_path, '<new1>')
+
+        self.pipe = pipe
+
+    def run(
+        self,
+        base_model: str,
+        weight_name: str,
+        prompt: str,
+        seed: int,
+        n_steps: int,
+        guidance_scale: float,
+        eta: float,
+        batch_size: int,
+        resolution: int,
+    ) -> PIL.Image.Image:
+        if not torch.cuda.is_available():
+            raise gr.Error('CUDA is not available.')
+
+        self.load_pipe(base_model, weight_name)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        out = self.pipe([prompt]*batch_size,
+                        num_inference_steps=n_steps,
+                        guidance_scale=guidance_scale,
+                        height=resolution, width=resolution,
+                        eta = eta,
+                        generator=generator)  # type: ignore
+        out = out.images
+        out = PIL.Image.fromarray(np.hstack([np.array(x) for x in out]))
+        return out

requirements.txt

@@ -0,0 +1,11 @@
+accelerate==0.15.0
+bitsandbytes==0.35.4
+diffusers==0.10.2
+ftfy==6.1.1
+Pillow==9.3.0
+torch==1.13.0
+torchvision==0.14.0
+transformers==4.25.1
+triton==2.0.0.dev20220701
+xformers==0.0.13
+clip_retrieval

style.css

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

trainer.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import os
+import pathlib
+import shlex
+import shutil
+import subprocess
+
+import gradio as gr
+import PIL.Image
+import torch
+
+os.environ['PYTHONPATH'] = f'custom-diffusion:{os.getenv("PYTHONPATH", "")}'
+
+
+def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
+    w, h = image.size
+    if w == h:
+        return image
+    elif w > h:
+        new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0))
+        new_image.paste(image, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0))
+        new_image.paste(image, ((h - w) // 2, 0))
+        return new_image
+
+
+class Trainer:
+    def __init__(self):
+        self.is_running = False
+        self.is_running_message = 'Another training is in progress.'
+
+        self.output_dir = pathlib.Path('results')
+        self.instance_data_dir = self.output_dir / 'training_data'
+        self.class_data_dir = self.output_dir / 'regularization_data'
+
+    def check_if_running(self) -> dict:
+        if self.is_running:
+            return gr.update(value=self.is_running_message)
+        else:
+            return gr.update(value='No training is running.')
+
+    def cleanup_dirs(self) -> None:
+        shutil.rmtree(self.output_dir, ignore_errors=True)
+
+    def prepare_dataset(self, concept_images: list, resolution: int) -> None:
+        self.instance_data_dir.mkdir(parents=True)
+        for i, temp_path in enumerate(concept_images):
+            image = PIL.Image.open(temp_path.name)
+            image = pad_image(image)
+            image = image.resize((resolution, resolution))
+            image = image.convert('RGB')
+            out_path = self.instance_data_dir / f'{i:03d}.jpg'
+            image.save(out_path, format='JPEG', quality=100)
+
+    def run(
+        self,
+        base_model: str,
+        resolution_s: str,
+        concept_images: list | None,
+        concept_prompt: str,
+        class_prompt: str,
+        n_steps: int,
+        learning_rate: float,
+        train_text_encoder: bool,
+        modifier_token: bool,
+        gradient_accumulation: int,
+        batch_size: int,
+        use_8bit_adam: bool,
+        gradient_checkpointing: bool,
+        gen_images: bool,
+    ) -> tuple[dict, list[pathlib.Path]]:
+        if not torch.cuda.is_available():
+            raise gr.Error('CUDA is not available.')
+
+        if self.is_running:
+            return gr.update(value=self.is_running_message), []
+
+        if concept_images is None:
+            raise gr.Error('You need to upload images.')
+        if not concept_prompt:
+            raise gr.Error('The concept prompt is missing.')
+
+        resolution = int(resolution_s)
+
+        self.cleanup_dirs()
+        self.prepare_dataset(concept_images, resolution)
+
+        command = f'''
+        accelerate launch custom-diffusion/src/diffuser_training.py \
+          --pretrained_model_name_or_path={base_model}   \
+          --instance_data_dir={self.instance_data_dir}  \
+          --output_dir={self.output_dir} \
+          --instance_prompt="{concept_prompt}" \
+          --class_data_dir={self.class_data_dir} \
+          --with_prior_preservation --prior_loss_weight=1.0 \
+          --class_prompt="{class_prompt}" \
+          --resolution={resolution}  \
+          --train_batch_size={batch_size}  \
+          --gradient_accumulation_steps={gradient_accumulation}  \
+          --learning_rate={learning_rate}  \
+          --lr_scheduler="constant" \
+          --lr_warmup_steps=0 \
+          --max_train_steps={n_steps} \
+          --num_class_images=200 \
+          --scale_lr 
+        '''
+        if modifier_token:
+            command += ' --modifier_token "<new1>"'
+        if not gen_images:
+            command += ' --real_prior'
+        if use_8bit_adam:
+            command += ' --use_8bit_adam'
+        if train_text_encoder:
+            command += f' --train_text_encoder'
+        if gradient_checkpointing:
+            command += f' --gradient_checkpointing'
+
+        with open(self.output_dir / 'train.sh', 'w') as f:
+            command_s = ' '.join(command.split())
+            f.write(command_s)
+
+        self.is_running = True
+        res = subprocess.run(shlex.split(command))
+        self.is_running = False
+
+        if res.returncode == 0:
+            result_message = 'Training Completed!'
+        else:
+            result_message = 'Training Failed!'
+        weight_paths = sorted(self.output_dir.glob('*.bin'))
+        return gr.update(value=result_message), weight_paths

uploader.py

@@ -0,0 +1,20 @@
+import gradio as gr
+from huggingface_hub import HfApi
+
+
+def upload(model_name: str, hf_token: str) -> None:
+    api = HfApi(token=hf_token)
+    user_name = api.whoami()['name']
+    model_id = f'{user_name}/{model_name}'
+    try:
+        api.create_repo(model_id, repo_type='model', private=True)
+        api.upload_folder(repo_id=model_id,
+                          folder_path='results',
+                          path_in_repo='results',
+                          repo_type='model')
+        url = f'https://huggingface.co/{model_id}'
+        message = f'Your model was successfully uploaded to [{url}]({url}).'
+    except Exception as e:
+        message = str(e)
+
+    return gr.update(value=message, visible=True)