hysts HF staff patrickvonplaten multimodalart HF staff commited on
Commit
5555256
0 Parent(s):

initial commit

Browse files

Co-authored-by: Patrick von Platen <patrickvonplaten@users.noreply.huggingface.co>
Co-authored-by: Apolinário from multimodal AI art <multimodalart@users.noreply.huggingface.co>

Files changed (10) hide show
  1. .gitattributes +35 -0
  2. .gitignore +162 -0
  3. .pre-commit-config.yaml +38 -0
  4. .style.yapf +5 -0
  5. README.md +14 -0
  6. app.py +648 -0
  7. model.py +309 -0
  8. requirements.txt +16 -0
  9. settings.py +56 -0
  10. style.css +203 -0
.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tflite filter=lfs diff=lfs merge=lfs -text
29
+ *.tgz filter=lfs diff=lfs merge=lfs -text
30
+ *.wasm filter=lfs diff=lfs merge=lfs -text
31
+ *.xz filter=lfs diff=lfs merge=lfs -text
32
+ *.zip filter=lfs diff=lfs merge=lfs -text
33
+ *.zst filter=lfs diff=lfs merge=lfs -text
34
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
35
+ *.png filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gradio_cached_examples/
2
+
3
+ # Byte-compiled / optimized / DLL files
4
+ __pycache__/
5
+ *.py[cod]
6
+ *$py.class
7
+
8
+ # C extensions
9
+ *.so
10
+
11
+ # Distribution / packaging
12
+ .Python
13
+ build/
14
+ develop-eggs/
15
+ dist/
16
+ downloads/
17
+ eggs/
18
+ .eggs/
19
+ lib/
20
+ lib64/
21
+ parts/
22
+ sdist/
23
+ var/
24
+ wheels/
25
+ share/python-wheels/
26
+ *.egg-info/
27
+ .installed.cfg
28
+ *.egg
29
+ MANIFEST
30
+
31
+ # PyInstaller
32
+ # Usually these files are written by a python script from a template
33
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
34
+ *.manifest
35
+ *.spec
36
+
37
+ # Installer logs
38
+ pip-log.txt
39
+ pip-delete-this-directory.txt
40
+
41
+ # Unit test / coverage reports
42
+ htmlcov/
43
+ .tox/
44
+ .nox/
45
+ .coverage
46
+ .coverage.*
47
+ .cache
48
+ nosetests.xml
49
+ coverage.xml
50
+ *.cover
51
+ *.py,cover
52
+ .hypothesis/
53
+ .pytest_cache/
54
+ cover/
55
+
56
+ # Translations
57
+ *.mo
58
+ *.pot
59
+
60
+ # Django stuff:
61
+ *.log
62
+ local_settings.py
63
+ db.sqlite3
64
+ db.sqlite3-journal
65
+
66
+ # Flask stuff:
67
+ instance/
68
+ .webassets-cache
69
+
70
+ # Scrapy stuff:
71
+ .scrapy
72
+
73
+ # Sphinx documentation
74
+ docs/_build/
75
+
76
+ # PyBuilder
77
+ .pybuilder/
78
+ target/
79
+
80
+ # Jupyter Notebook
81
+ .ipynb_checkpoints
82
+
83
+ # IPython
84
+ profile_default/
85
+ ipython_config.py
86
+
87
+ # pyenv
88
+ # For a library or package, you might want to ignore these files since the code is
89
+ # intended to run in multiple environments; otherwise, check them in:
90
+ # .python-version
91
+
92
+ # pipenv
93
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
94
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
95
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
96
+ # install all needed dependencies.
97
+ #Pipfile.lock
98
+
99
+ # poetry
100
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
101
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
102
+ # commonly ignored for libraries.
103
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
104
+ #poetry.lock
105
+
106
+ # pdm
107
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
108
+ #pdm.lock
109
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
110
+ # in version control.
111
+ # https://pdm.fming.dev/#use-with-ide
112
+ .pdm.toml
113
+
114
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
115
+ __pypackages__/
116
+
117
+ # Celery stuff
118
+ celerybeat-schedule
119
+ celerybeat.pid
120
+
121
+ # SageMath parsed files
122
+ *.sage.py
123
+
124
+ # Environments
125
+ .env
126
+ .venv
127
+ env/
128
+ venv/
129
+ ENV/
130
+ env.bak/
131
+ venv.bak/
132
+
133
+ # Spyder project settings
134
+ .spyderproject
135
+ .spyproject
136
+
137
+ # Rope project settings
138
+ .ropeproject
139
+
140
+ # mkdocs documentation
141
+ /site
142
+
143
+ # mypy
144
+ .mypy_cache/
145
+ .dmypy.json
146
+ dmypy.json
147
+
148
+ # Pyre type checker
149
+ .pyre/
150
+
151
+ # pytype static type analyzer
152
+ .pytype/
153
+
154
+ # Cython debug symbols
155
+ cython_debug/
156
+
157
+ # PyCharm
158
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
159
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
160
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
161
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
162
+ #.idea/
.pre-commit-config.yaml ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ exclude: diffusers-private
2
+ repos:
3
+ - repo: https://github.com/pre-commit/pre-commit-hooks
4
+ rev: v4.2.0
5
+ hooks:
6
+ - id: check-executables-have-shebangs
7
+ - id: check-json
8
+ - id: check-merge-conflict
9
+ - id: check-shebang-scripts-are-executable
10
+ - id: check-toml
11
+ - id: check-yaml
12
+ - id: double-quote-string-fixer
13
+ - id: end-of-file-fixer
14
+ - id: mixed-line-ending
15
+ args: ['--fix=lf']
16
+ - id: requirements-txt-fixer
17
+ - id: trailing-whitespace
18
+ - repo: https://github.com/myint/docformatter
19
+ rev: v1.4
20
+ hooks:
21
+ - id: docformatter
22
+ args: ['--in-place']
23
+ - repo: https://github.com/pycqa/isort
24
+ rev: 5.12.0
25
+ hooks:
26
+ - id: isort
27
+ - repo: https://github.com/pre-commit/mirrors-mypy
28
+ rev: v0.991
29
+ hooks:
30
+ - id: mypy
31
+ args: ['--ignore-missing-imports']
32
+ additional_dependencies: ['types-python-slugify']
33
+ files: ^diffusers-private
34
+ - repo: https://github.com/google/yapf
35
+ rev: v0.32.0
36
+ hooks:
37
+ - id: yapf
38
+ args: ['--parallel', '--in-place']
.style.yapf ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
1
+ [style]
2
+ based_on_style = pep8
3
+ blank_line_before_nested_class_or_def = false
4
+ spaces_before_comment = 2
5
+ split_before_logical_operator = true
README.md ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: IF
3
+ emoji: 🔥
4
+ colorFrom: pink
5
+ colorTo: red
6
+ sdk: gradio
7
+ sdk_version: 3.27.0
8
+ python_version: 3.10.11
9
+ app_file: app.py
10
+ pinned: false
11
+ license: other
12
+ ---
13
+
14
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,648 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ import datetime
4
+ import hashlib
5
+ import json
6
+ import os
7
+ import random
8
+ import tempfile
9
+
10
+ import gradio as gr
11
+ import torch
12
+ from huggingface_hub import HfApi
13
+
14
+ # isort: off
15
+ from model import Model
16
+ from settings import (
17
+ DEBUG,
18
+ DEFAULT_CUSTOM_TIMESTEPS_1,
19
+ DEFAULT_CUSTOM_TIMESTEPS_2,
20
+ DEFAULT_NUM_IMAGES,
21
+ DEFAULT_NUM_STEPS_3,
22
+ DISABLE_SD_X4_UPSCALER,
23
+ GALLERY_COLUMN_NUM,
24
+ HF_TOKEN,
25
+ MAX_NUM_IMAGES,
26
+ MAX_NUM_STEPS,
27
+ MAX_QUEUE_SIZE,
28
+ MAX_SEED,
29
+ SHOW_ADVANCED_OPTIONS,
30
+ SHOW_CUSTOM_TIMESTEPS_1,
31
+ SHOW_CUSTOM_TIMESTEPS_2,
32
+ SHOW_DEVICE_WARNING,
33
+ SHOW_DUPLICATE_BUTTON,
34
+ SHOW_NUM_IMAGES,
35
+ SHOW_NUM_STEPS_1,
36
+ SHOW_NUM_STEPS_2,
37
+ SHOW_NUM_STEPS_3,
38
+ SHOW_UPSCALE_TO_256_BUTTON,
39
+ UPLOAD_REPO_ID,
40
+ UPLOAD_RESULT_IMAGE,
41
+ )
42
+ # isort: on
43
+
44
+ DESCRIPTION = '# [DeepFloyd IF](https://github.com/deep-floyd/IF)'
45
+
46
+ if SHOW_DUPLICATE_BUTTON:
47
+ SPACE_ID = os.getenv('SPACE_ID')
48
+ DESCRIPTION += f'\n<p><a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space%20to%20skip%20the%20queue-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></p>'
49
+
50
+ if SHOW_DEVICE_WARNING and not torch.cuda.is_available():
51
+ DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'
52
+
53
+ model = Model()
54
+
55
+
56
+ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
57
+ if randomize_seed:
58
+ seed = random.randint(0, MAX_SEED)
59
+ return seed
60
+
61
+
62
+ def get_stage2_index(evt: gr.SelectData) -> int:
63
+ return evt.index
64
+
65
+
66
+ def check_if_stage2_selected(index: int) -> None:
67
+ if index == -1:
68
+ raise gr.Error(
69
+ 'You need to select the image you would like to upscale from the Stage 1 results by clicking.'
70
+ )
71
+
72
+
73
+ hf_api = HfApi(token=HF_TOKEN)
74
+ if UPLOAD_REPO_ID:
75
+ hf_api.create_repo(repo_id=UPLOAD_REPO_ID,
76
+ private=True,
77
+ repo_type='dataset',
78
+ exist_ok=True)
79
+
80
+
81
+ def get_param_file_hash_name(param_filepath: str) -> str:
82
+ if not UPLOAD_REPO_ID:
83
+ return ''
84
+ with open(param_filepath, 'rb') as f:
85
+ md5 = hashlib.md5(f.read()).hexdigest()
86
+ utcnow = datetime.datetime.utcnow().strftime('%Y-%m-%d-%H-%M-%S-%f')
87
+ return f'{utcnow}-{md5}'
88
+
89
+
90
+ def upload_stage1_result(stage1_param_path: str, stage1_result_path: str,
91
+ save_name: str) -> None:
92
+ if not UPLOAD_REPO_ID:
93
+ return
94
+ try:
95
+ hf_api.upload_file(path_or_fileobj=stage1_param_path,
96
+ path_in_repo=f'stage1_params/{save_name}.json',
97
+ repo_id=UPLOAD_REPO_ID,
98
+ repo_type='dataset')
99
+ hf_api.upload_file(path_or_fileobj=stage1_result_path,
100
+ path_in_repo=f'stage1_results/{save_name}.pth',
101
+ repo_id=UPLOAD_REPO_ID,
102
+ repo_type='dataset')
103
+ except Exception as e:
104
+ print(e)
105
+
106
+
107
+ def upload_stage2_info(stage1_param_file_hash_name: str,
108
+ stage2_output_path: str,
109
+ selected_index_for_upscale: int, seed_2: int,
110
+ guidance_scale_2: float, custom_timesteps_2: str,
111
+ num_inference_steps_2: int) -> None:
112
+ if not UPLOAD_REPO_ID:
113
+ return
114
+ if not stage1_param_file_hash_name:
115
+ raise ValueError
116
+
117
+ stage2_params = {
118
+ 'stage1_param_file_hash_name': stage1_param_file_hash_name,
119
+ 'selected_index_for_upscale': selected_index_for_upscale,
120
+ 'seed_2': seed_2,
121
+ 'guidance_scale_2': guidance_scale_2,
122
+ 'custom_timesteps_2': custom_timesteps_2,
123
+ 'num_inference_steps_2': num_inference_steps_2,
124
+ }
125
+ with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
126
+ param_file.write(json.dumps(stage2_params))
127
+ stage2_param_file_hash_name = get_param_file_hash_name(param_file.name)
128
+ save_name = f'{stage1_param_file_hash_name}_{stage2_param_file_hash_name}'
129
+
130
+ try:
131
+ hf_api.upload_file(path_or_fileobj=param_file.name,
132
+ path_in_repo=f'stage2_params/{save_name}.json',
133
+ repo_id=UPLOAD_REPO_ID,
134
+ repo_type='dataset')
135
+ if UPLOAD_RESULT_IMAGE:
136
+ hf_api.upload_file(path_or_fileobj=stage2_output_path,
137
+ path_in_repo=f'stage2_results/{save_name}.png',
138
+ repo_id=UPLOAD_REPO_ID,
139
+ repo_type='dataset')
140
+ except Exception as e:
141
+ print(e)
142
+
143
+
144
+ def upload_stage2_3_info(stage1_param_file_hash_name: str,
145
+ stage2_3_output_path: str,
146
+ selected_index_for_upscale: int, seed_2: int,
147
+ guidance_scale_2: float, custom_timesteps_2: str,
148
+ num_inference_steps_2: int, prompt: str,
149
+ negative_prompt: str, seed_3: int,
150
+ guidance_scale_3: float,
151
+ num_inference_steps_3: int) -> None:
152
+ if not UPLOAD_REPO_ID:
153
+ return
154
+ if not stage1_param_file_hash_name:
155
+ raise ValueError
156
+
157
+ stage2_3_params = {
158
+ 'stage1_param_file_hash_name': stage1_param_file_hash_name,
159
+ 'selected_index_for_upscale': selected_index_for_upscale,
160
+ 'seed_2': seed_2,
161
+ 'guidance_scale_2': guidance_scale_2,
162
+ 'custom_timesteps_2': custom_timesteps_2,
163
+ 'num_inference_steps_2': num_inference_steps_2,
164
+ 'prompt': prompt,
165
+ 'negative_prompt': negative_prompt,
166
+ 'seed_3': seed_3,
167
+ 'guidance_scale_3': guidance_scale_3,
168
+ 'num_inference_steps_3': num_inference_steps_3,
169
+ }
170
+ with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
171
+ param_file.write(json.dumps(stage2_3_params))
172
+ stage2_3_param_file_hash_name = get_param_file_hash_name(param_file.name)
173
+ save_name = f'{stage1_param_file_hash_name}_{stage2_3_param_file_hash_name}'
174
+
175
+ try:
176
+ hf_api.upload_file(path_or_fileobj=param_file.name,
177
+ path_in_repo=f'stage2_3_params/{save_name}.json',
178
+ repo_id=UPLOAD_REPO_ID,
179
+ repo_type='dataset')
180
+ if UPLOAD_RESULT_IMAGE:
181
+ hf_api.upload_file(
182
+ path_or_fileobj=stage2_3_output_path,
183
+ path_in_repo=f'stage2_3_results/{save_name}.png',
184
+ repo_id=UPLOAD_REPO_ID,
185
+ repo_type='dataset')
186
+ except Exception as e:
187
+ print(e)
188
+
189
+
190
+ def update_upscale_button(selected_index: int) -> tuple[dict, dict]:
191
+ if selected_index == -1:
192
+ return gr.update(interactive=False), gr.update(interactive=False)
193
+ else:
194
+ return gr.update(interactive=True), gr.update(interactive=True)
195
+
196
+
197
+ def _update_result_view(show_gallery: bool) -> tuple[dict, dict]:
198
+ return gr.update(visible=show_gallery), gr.update(visible=not show_gallery)
199
+
200
+
201
+ def show_gallery_view() -> tuple[dict, dict]:
202
+ return _update_result_view(True)
203
+
204
+
205
+ def show_upscaled_view() -> tuple[dict, dict]:
206
+ return _update_result_view(False)
207
+
208
+
209
+ examples = [
210
+ 'high quality dslr photo, a photo product of a lemon inspired by natural and organic materials, wooden accents, intricately decorated with glowing vines of led lights, inspired by baroque luxury',
211
+ 'Aerial photo of a beach, the words "what if?" written in the sand.',
212
+ 'A photo of a red cube on top of a blue cube. a photo of a red cube with text "blue" on it is sitting on top of a blue cube with text "red" on it. photo realism',
213
+ 'a photo of a violet baseball cap with yellow text: "deep floyd". 50mm lens, photo realism, cine lens. violet baseball cap says "deep floyd". reflections, render. yellow stitch text "deep floyd"',
214
+ 'ultra close-up color photo portrait of rainbow owl with deer horns in the woods',
215
+ 'product image of a crochet Cthulhu the great old one emerging from a spacetime wormhole made of wool.',
216
+ 'a little green budgie parrot driving small red toy car in new york street, photo',
217
+ 'origami dancer in white paper, 3d render, ultra-detailed, on white background, studio shot.',
218
+ 'glowing mushrooms in a natural environment with smoke in the frame',
219
+ 'a bowl full of few adorable golden doodle puppies, the doodles dusted in powdered sugar and look delicious, bokeh, cannon. professional macro photo, super detailed. cute sweet golden doodle confectionery, baking puppies in powdered sugar in the bowl',
220
+ 'a yellow ipe tree in the cerrado, the hills are dry and the weather is hot. the ipe tree shows all its beauty among the dry trees. cinematic film still of a movie, realism, 4k, 8mm, grainy, panavision',
221
+ ]
222
+
223
+ with gr.Blocks(css='style.css') as demo:
224
+ gr.Markdown(DESCRIPTION)
225
+ with gr.Box():
226
+ with gr.Row(elem_id='prompt-container').style(equal_height=True):
227
+ with gr.Column():
228
+ prompt = gr.Text(
229
+ label='Prompt',
230
+ show_label=False,
231
+ max_lines=1,
232
+ placeholder='Enter your prompt',
233
+ elem_id='prompt-text-input',
234
+ ).style(container=False)
235
+ negative_prompt = gr.Text(
236
+ label='Negative prompt',
237
+ show_label=False,
238
+ max_lines=1,
239
+ placeholder='Enter a negative prompt',
240
+ elem_id='negative-prompt-text-input',
241
+ ).style(container=False)
242
+ generate_button = gr.Button('Generate').style(full_width=False)
243
+
244
+ with gr.Column() as gallery_view:
245
+ gallery = gr.Gallery(label='Stage 1 results',
246
+ show_label=False,
247
+ elem_id='gallery').style(
248
+ columns=GALLERY_COLUMN_NUM,
249
+ object_fit='contain')
250
+ gr.Markdown('Pick your favorite generation to upscale.')
251
+ with gr.Row():
252
+ upscale_to_256_button = gr.Button(
253
+ 'Upscale to 256px',
254
+ visible=SHOW_UPSCALE_TO_256_BUTTON
255
+ or DISABLE_SD_X4_UPSCALER,
256
+ interactive=False)
257
+ upscale_button = gr.Button('Upscale',
258
+ interactive=False,
259
+ visible=not DISABLE_SD_X4_UPSCALER)
260
+ with gr.Column(visible=False) as upscale_view:
261
+ result = gr.Image(label='Result',
262
+ show_label=False,
263
+ type='filepath',
264
+ interactive=False,
265
+ elem_id='upscaled-image').style(height=640)
266
+ back_to_selection_button = gr.Button('Back to selection')
267
+
268
+ with gr.Accordion('Advanced options',
269
+ open=False,
270
+ visible=SHOW_ADVANCED_OPTIONS):
271
+ with gr.Tabs():
272
+ with gr.Tab(label='Generation'):
273
+ seed_1 = gr.Slider(label='Seed',
274
+ minimum=0,
275
+ maximum=MAX_SEED,
276
+ step=1,
277
+ value=0)
278
+ randomize_seed_1 = gr.Checkbox(label='Randomize seed',
279
+ value=True)
280
+ guidance_scale_1 = gr.Slider(label='Guidance scale',
281
+ minimum=1,
282
+ maximum=20,
283
+ step=0.1,
284
+ value=7.0)
285
+ custom_timesteps_1 = gr.Dropdown(
286
+ label='Custom timesteps 1',
287
+ choices=[
288
+ 'none',
289
+ 'fast27',
290
+ 'smart27',
291
+ 'smart50',
292
+ 'smart100',
293
+ 'smart185',
294
+ ],
295
+ value=DEFAULT_CUSTOM_TIMESTEPS_1,
296
+ visible=SHOW_CUSTOM_TIMESTEPS_1)
297
+ num_inference_steps_1 = gr.Slider(
298
+ label='Number of inference steps',
299
+ minimum=1,
300
+ maximum=MAX_NUM_STEPS,
301
+ step=1,
302
+ value=100,
303
+ visible=SHOW_NUM_STEPS_1)
304
+ num_images = gr.Slider(label='Number of images',
305
+ minimum=1,
306
+ maximum=MAX_NUM_IMAGES,
307
+ step=1,
308
+ value=DEFAULT_NUM_IMAGES,
309
+ visible=SHOW_NUM_IMAGES)
310
+ with gr.Tab(label='Super-resolution 1'):
311
+ seed_2 = gr.Slider(label='Seed',
312
+ minimum=0,
313
+ maximum=MAX_SEED,
314
+ step=1,
315
+ value=0)
316
+ randomize_seed_2 = gr.Checkbox(label='Randomize seed',
317
+ value=True)
318
+ guidance_scale_2 = gr.Slider(label='Guidance scale',
319
+ minimum=1,
320
+ maximum=20,
321
+ step=0.1,
322
+ value=4.0)
323
+ custom_timesteps_2 = gr.Dropdown(
324
+ label='Custom timesteps 2',
325
+ choices=[
326
+ 'none',
327
+ 'fast27',
328
+ 'smart27',
329
+ 'smart50',
330
+ 'smart100',
331
+ 'smart185',
332
+ ],
333
+ value=DEFAULT_CUSTOM_TIMESTEPS_2,
334
+ visible=SHOW_CUSTOM_TIMESTEPS_2)
335
+ num_inference_steps_2 = gr.Slider(
336
+ label='Number of inference steps',
337
+ minimum=1,
338
+ maximum=MAX_NUM_STEPS,
339
+ step=1,
340
+ value=50,
341
+ visible=SHOW_NUM_STEPS_2)
342
+ with gr.Tab(label='Super-resolution 2'):
343
+ seed_3 = gr.Slider(label='Seed',
344
+ minimum=0,
345
+ maximum=MAX_SEED,
346
+ step=1,
347
+ value=0)
348
+ randomize_seed_3 = gr.Checkbox(label='Randomize seed',
349
+ value=True)
350
+ guidance_scale_3 = gr.Slider(label='Guidance scale',
351
+ minimum=1,
352
+ maximum=20,
353
+ step=0.1,
354
+ value=9.0)
355
+ num_inference_steps_3 = gr.Slider(
356
+ label='Number of inference steps',
357
+ minimum=1,
358
+ maximum=MAX_NUM_STEPS,
359
+ step=1,
360
+ value=DEFAULT_NUM_STEPS_3,
361
+ visible=SHOW_NUM_STEPS_3)
362
+
363
+ gr.Examples(examples=examples, inputs=prompt, examples_per_page=4)
364
+
365
+ with gr.Box(visible=DEBUG):
366
+ with gr.Row():
367
+ with gr.Accordion(label='Hidden params'):
368
+ stage1_param_path = gr.Text(label='Stage 1 param path')
369
+ stage1_result_path = gr.Text(label='Stage 1 result path')
370
+ stage1_param_file_hash_name = gr.Text(
371
+ label='Stage 1 param file hash name')
372
+ selected_index_for_stage2 = gr.Number(
373
+ label='Selected index for Stage 2', value=-1, precision=0)
374
+
375
+ stage1_inputs = [
376
+ prompt,
377
+ negative_prompt,
378
+ seed_1,
379
+ num_images,
380
+ guidance_scale_1,
381
+ custom_timesteps_1,
382
+ num_inference_steps_1,
383
+ ]
384
+ stage1_outputs = [
385
+ gallery,
386
+ stage1_param_path,
387
+ stage1_result_path,
388
+ ]
389
+
390
+ prompt.submit(
391
+ fn=randomize_seed_fn,
392
+ inputs=[seed_1, randomize_seed_1],
393
+ outputs=seed_1,
394
+ queue=False,
395
+ ).then(
396
+ fn=lambda: -1,
397
+ outputs=selected_index_for_stage2,
398
+ queue=False,
399
+ ).then(
400
+ fn=show_gallery_view,
401
+ outputs=[
402
+ gallery_view,
403
+ upscale_view,
404
+ ],
405
+ queue=False,
406
+ ).then(
407
+ fn=update_upscale_button,
408
+ inputs=selected_index_for_stage2,
409
+ outputs=[
410
+ upscale_button,
411
+ upscale_to_256_button,
412
+ ],
413
+ queue=False,
414
+ ).then(
415
+ fn=model.run_stage1,
416
+ inputs=stage1_inputs,
417
+ outputs=stage1_outputs,
418
+ ).success(
419
+ fn=get_param_file_hash_name,
420
+ inputs=stage1_param_path,
421
+ outputs=stage1_param_file_hash_name,
422
+ queue=False,
423
+ ).then(
424
+ fn=upload_stage1_result,
425
+ inputs=[
426
+ stage1_param_path,
427
+ stage1_result_path,
428
+ stage1_param_file_hash_name,
429
+ ],
430
+ queue=False,
431
+ )
432
+
433
+ negative_prompt.submit(
434
+ fn=randomize_seed_fn,
435
+ inputs=[seed_1, randomize_seed_1],
436
+ outputs=seed_1,
437
+ queue=False,
438
+ ).then(
439
+ fn=lambda: -1,
440
+ outputs=selected_index_for_stage2,
441
+ queue=False,
442
+ ).then(
443
+ fn=show_gallery_view,
444
+ outputs=[
445
+ gallery_view,
446
+ upscale_view,
447
+ ],
448
+ queue=False,
449
+ ).then(
450
+ fn=update_upscale_button,
451
+ inputs=selected_index_for_stage2,
452
+ outputs=[
453
+ upscale_button,
454
+ upscale_to_256_button,
455
+ ],
456
+ queue=False,
457
+ ).then(
458
+ fn=model.run_stage1,
459
+ inputs=stage1_inputs,
460
+ outputs=stage1_outputs,
461
+ ).success(
462
+ fn=get_param_file_hash_name,
463
+ inputs=stage1_param_path,
464
+ outputs=stage1_param_file_hash_name,
465
+ queue=False,
466
+ ).then(
467
+ fn=upload_stage1_result,
468
+ inputs=[
469
+ stage1_param_path,
470
+ stage1_result_path,
471
+ stage1_param_file_hash_name,
472
+ ],
473
+ queue=False,
474
+ )
475
+
476
+ generate_button.click(
477
+ fn=randomize_seed_fn,
478
+ inputs=[seed_1, randomize_seed_1],
479
+ outputs=seed_1,
480
+ queue=False,
481
+ ).then(
482
+ fn=lambda: -1,
483
+ outputs=selected_index_for_stage2,
484
+ queue=False,
485
+ ).then(
486
+ fn=show_gallery_view,
487
+ outputs=[
488
+ gallery_view,
489
+ upscale_view,
490
+ ],
491
+ queue=False,
492
+ ).then(
493
+ fn=update_upscale_button,
494
+ inputs=selected_index_for_stage2,
495
+ outputs=[
496
+ upscale_button,
497
+ upscale_to_256_button,
498
+ ],
499
+ queue=False,
500
+ ).then(
501
+ fn=model.run_stage1,
502
+ inputs=stage1_inputs,
503
+ outputs=stage1_outputs,
504
+ api_name='generate64',
505
+ ).success(
506
+ fn=get_param_file_hash_name,
507
+ inputs=stage1_param_path,
508
+ outputs=stage1_param_file_hash_name,
509
+ queue=False,
510
+ ).then(
511
+ fn=upload_stage1_result,
512
+ inputs=[
513
+ stage1_param_path,
514
+ stage1_result_path,
515
+ stage1_param_file_hash_name,
516
+ ],
517
+ queue=False,
518
+ )
519
+
520
+ gallery.select(
521
+ fn=get_stage2_index,
522
+ outputs=selected_index_for_stage2,
523
+ queue=False,
524
+ )
525
+
526
+ selected_index_for_stage2.change(
527
+ fn=update_upscale_button,
528
+ inputs=selected_index_for_stage2,
529
+ outputs=[
530
+ upscale_button,
531
+ upscale_to_256_button,
532
+ ],
533
+ queue=False,
534
+ )
535
+
536
+ stage2_inputs = [
537
+ stage1_result_path,
538
+ selected_index_for_stage2,
539
+ seed_2,
540
+ guidance_scale_2,
541
+ custom_timesteps_2,
542
+ num_inference_steps_2,
543
+ ]
544
+
545
+ upscale_to_256_button.click(
546
+ fn=check_if_stage2_selected,
547
+ inputs=selected_index_for_stage2,
548
+ queue=False,
549
+ ).then(
550
+ fn=randomize_seed_fn,
551
+ inputs=[seed_2, randomize_seed_2],
552
+ outputs=seed_2,
553
+ queue=False,
554
+ ).then(
555
+ fn=show_upscaled_view,
556
+ outputs=[
557
+ gallery_view,
558
+ upscale_view,
559
+ ],
560
+ queue=False,
561
+ ).then(
562
+ fn=model.run_stage2,
563
+ inputs=stage2_inputs,
564
+ outputs=result,
565
+ api_name='upscale256',
566
+ ).success(
567
+ fn=upload_stage2_info,
568
+ inputs=[
569
+ stage1_param_file_hash_name,
570
+ result,
571
+ selected_index_for_stage2,
572
+ seed_2,
573
+ guidance_scale_2,
574
+ custom_timesteps_2,
575
+ num_inference_steps_2,
576
+ ],
577
+ queue=False,
578
+ )
579
+
580
+ stage2_3_inputs = [
581
+ stage1_result_path,
582
+ selected_index_for_stage2,
583
+ seed_2,
584
+ guidance_scale_2,
585
+ custom_timesteps_2,
586
+ num_inference_steps_2,
587
+ prompt,
588
+ negative_prompt,
589
+ seed_3,
590
+ guidance_scale_3,
591
+ num_inference_steps_3,
592
+ ]
593
+
594
+ upscale_button.click(
595
+ fn=check_if_stage2_selected,
596
+ inputs=selected_index_for_stage2,
597
+ queue=False,
598
+ ).then(
599
+ fn=randomize_seed_fn,
600
+ inputs=[seed_2, randomize_seed_2],
601
+ outputs=seed_2,
602
+ queue=False,
603
+ ).then(
604
+ fn=randomize_seed_fn,
605
+ inputs=[seed_3, randomize_seed_3],
606
+ outputs=seed_3,
607
+ queue=False,
608
+ ).then(
609
+ fn=show_upscaled_view,
610
+ outputs=[
611
+ gallery_view,
612
+ upscale_view,
613
+ ],
614
+ queue=False,
615
+ ).then(
616
+ fn=model.run_stage2_3,
617
+ inputs=stage2_3_inputs,
618
+ outputs=result,
619
+ api_name='upscale1024',
620
+ ).success(
621
+ fn=upload_stage2_3_info,
622
+ inputs=[
623
+ stage1_param_file_hash_name,
624
+ result,
625
+ selected_index_for_stage2,
626
+ seed_2,
627
+ guidance_scale_2,
628
+ custom_timesteps_2,
629
+ num_inference_steps_2,
630
+ prompt,
631
+ negative_prompt,
632
+ seed_3,
633
+ guidance_scale_3,
634
+ num_inference_steps_3,
635
+ ],
636
+ queue=False,
637
+ )
638
+
639
+ back_to_selection_button.click(
640
+ fn=show_gallery_view,
641
+ outputs=[
642
+ gallery_view,
643
+ upscale_view,
644
+ ],
645
+ queue=False,
646
+ )
647
+
648
+ demo.queue(api_open=False, max_size=MAX_QUEUE_SIZE).launch(debug=DEBUG)
model.py ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import gc
4
+ import json
5
+ import tempfile
6
+ from typing import Generator
7
+
8
+ import numpy as np
9
+ import PIL.Image
10
+ import torch
11
+ from diffusers import DiffusionPipeline, StableDiffusionUpscalePipeline
12
+ from diffusers.pipelines.deepfloyd_if import (fast27_timesteps,
13
+ smart27_timesteps,
14
+ smart50_timesteps,
15
+ smart100_timesteps,
16
+ smart185_timesteps)
17
+
18
+ from settings import (DISABLE_AUTOMATIC_CPU_OFFLOAD, DISABLE_SD_X4_UPSCALER,
19
+ HF_TOKEN, MAX_NUM_IMAGES, MAX_NUM_STEPS, MAX_SEED,
20
+ RUN_GARBAGE_COLLECTION)
21
+
22
+
23
+ class Model:
24
+ def __init__(self):
25
+ self.device = torch.device(
26
+ 'cuda:0' if torch.cuda.is_available() else 'cpu')
27
+ self.pipe = None
28
+ self.super_res_1_pipe = None
29
+ self.super_res_2_pipe = None
30
+ self.watermark_image = None
31
+
32
+ if torch.cuda.is_available():
33
+ self.load_weights()
34
+ self.watermark_image = PIL.Image.fromarray(
35
+ self.pipe.watermarker.watermark_image.to(
36
+ torch.uint8).cpu().numpy(),
37
+ mode='RGBA')
38
+
39
+ def load_weights(self) -> None:
40
+ self.pipe = DiffusionPipeline.from_pretrained(
41
+ 'DeepFloyd/IF-I-IF-v1.0',
42
+ torch_dtype=torch.float16,
43
+ variant='fp16',
44
+ use_safetensors=True,
45
+ use_auth_token=HF_TOKEN)
46
+ self.super_res_1_pipe = DiffusionPipeline.from_pretrained(
47
+ 'DeepFloyd/IF-II-L-v1.0',
48
+ text_encoder=None,
49
+ torch_dtype=torch.float16,
50
+ variant='fp16',
51
+ use_safetensors=True,
52
+ use_auth_token=HF_TOKEN)
53
+
54
+ if not DISABLE_SD_X4_UPSCALER:
55
+ self.super_res_2_pipe = StableDiffusionUpscalePipeline.from_pretrained(
56
+ 'stabilityai/stable-diffusion-x4-upscaler',
57
+ torch_dtype=torch.float16)
58
+
59
+ if DISABLE_AUTOMATIC_CPU_OFFLOAD:
60
+ self.pipe.to(self.device)
61
+ self.super_res_1_pipe.to(self.device)
62
+ if not DISABLE_SD_X4_UPSCALER:
63
+ self.super_res_2_pipe.to(self.device)
64
+ else:
65
+ self.pipe.enable_model_cpu_offload()
66
+ self.super_res_1_pipe.enable_model_cpu_offload()
67
+ if not DISABLE_SD_X4_UPSCALER:
68
+ self.super_res_2_pipe.enable_model_cpu_offload()
69
+
70
+ def apply_watermark_to_sd_x4_upscaler_results(
71
+ self, images: list[PIL.Image.Image]) -> None:
72
+ w, h = images[0].size
73
+
74
+ stability_x4_upscaler_sample_size = 128
75
+
76
+ coef = min(h / stability_x4_upscaler_sample_size,
77
+ w / stability_x4_upscaler_sample_size)
78
+ img_h, img_w = (int(h / coef), int(w / coef)) if coef < 1 else (h, w)
79
+
80
+ S1, S2 = 1024**2, img_w * img_h
81
+ K = (S2 / S1)**0.5
82
+ watermark_size = int(K * 62)
83
+ watermark_x = img_w - int(14 * K)
84
+ watermark_y = img_h - int(14 * K)
85
+
86
+ watermark_image = self.watermark_image.copy().resize(
87
+ (watermark_size, watermark_size),
88
+ PIL.Image.Resampling.BICUBIC,
89
+ reducing_gap=None)
90
+
91
+ for image in images:
92
+ image.paste(watermark_image,
93
+ box=(
94
+ watermark_x - watermark_size,
95
+ watermark_y - watermark_size,
96
+ watermark_x,
97
+ watermark_y,
98
+ ),
99
+ mask=watermark_image.split()[-1])
100
+
101
+ @staticmethod
102
+ def to_pil_images(images: torch.Tensor) -> list[PIL.Image.Image]:
103
+ images = (images / 2 + 0.5).clamp(0, 1)
104
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
105
+ images = np.round(images * 255).astype(np.uint8)
106
+ return [PIL.Image.fromarray(image) for image in images]
107
+
108
+ @staticmethod
109
+ def check_seed(seed: int) -> None:
110
+ if not 0 <= seed <= MAX_SEED:
111
+ raise ValueError
112
+
113
+ @staticmethod
114
+ def check_num_images(num_images: int) -> None:
115
+ if not 1 <= num_images <= MAX_NUM_IMAGES:
116
+ raise ValueError
117
+
118
+ @staticmethod
119
+ def check_num_inference_steps(num_steps: int) -> None:
120
+ if not 1 <= num_steps <= MAX_NUM_STEPS:
121
+ raise ValueError
122
+
123
+ @staticmethod
124
+ def get_custom_timesteps(name: str) -> list[int] | None:
125
+ if name == 'none':
126
+ timesteps = None
127
+ elif name == 'fast27':
128
+ timesteps = fast27_timesteps
129
+ elif name == 'smart27':
130
+ timesteps = smart27_timesteps
131
+ elif name == 'smart50':
132
+ timesteps = smart50_timesteps
133
+ elif name == 'smart100':
134
+ timesteps = smart100_timesteps
135
+ elif name == 'smart185':
136
+ timesteps = smart185_timesteps
137
+ else:
138
+ raise ValueError
139
+ return timesteps
140
+
141
+ @staticmethod
142
+ def run_garbage_collection():
143
+ gc.collect()
144
+ torch.cuda.empty_cache()
145
+
146
+ def run_stage1(
147
+ self,
148
+ prompt: str,
149
+ negative_prompt: str = '',
150
+ seed: int = 0,
151
+ num_images: int = 1,
152
+ guidance_scale_1: float = 7.0,
153
+ custom_timesteps_1: str = 'smart100',
154
+ num_inference_steps_1: int = 100,
155
+ ) -> tuple[list[PIL.Image.Image], str, str]:
156
+ self.check_seed(seed)
157
+ self.check_num_images(num_images)
158
+ self.check_num_inference_steps(num_inference_steps_1)
159
+
160
+ if RUN_GARBAGE_COLLECTION:
161
+ self.run_garbage_collection()
162
+
163
+ generator = torch.Generator(device=self.device).manual_seed(seed)
164
+
165
+ prompt_embeds, negative_embeds = self.pipe.encode_prompt(
166
+ prompt=prompt, negative_prompt=negative_prompt)
167
+
168
+ timesteps = self.get_custom_timesteps(custom_timesteps_1)
169
+
170
+ images = self.pipe(prompt_embeds=prompt_embeds,
171
+ negative_prompt_embeds=negative_embeds,
172
+ num_images_per_prompt=num_images,
173
+ guidance_scale=guidance_scale_1,
174
+ timesteps=timesteps,
175
+ num_inference_steps=num_inference_steps_1,
176
+ generator=generator,
177
+ output_type='pt').images
178
+ pil_images = self.to_pil_images(images)
179
+ self.pipe.watermarker.apply_watermark(
180
+ pil_images, self.pipe.unet.config.sample_size)
181
+
182
+ stage1_params = {
183
+ 'prompt': prompt,
184
+ 'negative_prompt': negative_prompt,
185
+ 'seed': seed,
186
+ 'num_images': num_images,
187
+ 'guidance_scale_1': guidance_scale_1,
188
+ 'custom_timesteps_1': custom_timesteps_1,
189
+ 'num_inference_steps_1': num_inference_steps_1,
190
+ }
191
+ with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
192
+ param_file.write(json.dumps(stage1_params))
193
+ stage1_result = {
194
+ 'prompt_embeds': prompt_embeds,
195
+ 'negative_embeds': negative_embeds,
196
+ 'images': images,
197
+ 'pil_images': pil_images,
198
+ }
199
+ with tempfile.NamedTemporaryFile(delete=False) as result_file:
200
+ torch.save(stage1_result, result_file.name)
201
+ return pil_images, param_file.name, result_file.name
202
+
203
+ def run_stage2(
204
+ self,
205
+ stage1_result_path: str,
206
+ stage2_index: int,
207
+ seed_2: int = 0,
208
+ guidance_scale_2: float = 4.0,
209
+ custom_timesteps_2: str = 'smart50',
210
+ num_inference_steps_2: int = 50,
211
+ disable_watermark: bool = False,
212
+ ) -> PIL.Image.Image:
213
+ self.check_seed(seed_2)
214
+ self.check_num_inference_steps(num_inference_steps_2)
215
+
216
+ if RUN_GARBAGE_COLLECTION:
217
+ self.run_garbage_collection()
218
+
219
+ generator = torch.Generator(device=self.device).manual_seed(seed_2)
220
+
221
+ stage1_result = torch.load(stage1_result_path)
222
+ prompt_embeds = stage1_result['prompt_embeds']
223
+ negative_embeds = stage1_result['negative_embeds']
224
+ images = stage1_result['images']
225
+ images = images[[stage2_index]]
226
+
227
+ timesteps = self.get_custom_timesteps(custom_timesteps_2)
228
+
229
+ out = self.super_res_1_pipe(image=images,
230
+ prompt_embeds=prompt_embeds,
231
+ negative_prompt_embeds=negative_embeds,
232
+ num_images_per_prompt=1,
233
+ guidance_scale=guidance_scale_2,
234
+ timesteps=timesteps,
235
+ num_inference_steps=num_inference_steps_2,
236
+ generator=generator,
237
+ output_type='pt',
238
+ noise_level=250).images
239
+ pil_images = self.to_pil_images(out)
240
+
241
+ if disable_watermark:
242
+ return pil_images[0]
243
+
244
+ self.super_res_1_pipe.watermarker.apply_watermark(
245
+ pil_images, self.super_res_1_pipe.unet.config.sample_size)
246
+ return pil_images[0]
247
+
248
+ def run_stage3(
249
+ self,
250
+ image: PIL.Image.Image,
251
+ prompt: str = '',
252
+ negative_prompt: str = '',
253
+ seed_3: int = 0,
254
+ guidance_scale_3: float = 9.0,
255
+ num_inference_steps_3: int = 75,
256
+ ) -> PIL.Image.Image:
257
+ self.check_seed(seed_3)
258
+ self.check_num_inference_steps(num_inference_steps_3)
259
+
260
+ if RUN_GARBAGE_COLLECTION:
261
+ self.run_garbage_collection()
262
+
263
+ generator = torch.Generator(device=self.device).manual_seed(seed_3)
264
+ out = self.super_res_2_pipe(image=image,
265
+ prompt=prompt,
266
+ negative_prompt=negative_prompt,
267
+ num_images_per_prompt=1,
268
+ guidance_scale=guidance_scale_3,
269
+ num_inference_steps=num_inference_steps_3,
270
+ generator=generator,
271
+ noise_level=100).images
272
+ self.apply_watermark_to_sd_x4_upscaler_results(out)
273
+ return out[0]
274
+
275
+ def run_stage2_3(
276
+ self,
277
+ stage1_result_path: str,
278
+ stage2_index: int,
279
+ seed_2: int = 0,
280
+ guidance_scale_2: float = 4.0,
281
+ custom_timesteps_2: str = 'smart50',
282
+ num_inference_steps_2: int = 50,
283
+ prompt: str = '',
284
+ negative_prompt: str = '',
285
+ seed_3: int = 0,
286
+ guidance_scale_3: float = 9.0,
287
+ num_inference_steps_3: int = 75,
288
+ ) -> Generator[PIL.Image.Image]:
289
+ self.check_seed(seed_3)
290
+ self.check_num_inference_steps(num_inference_steps_3)
291
+
292
+ out_image = self.run_stage2(
293
+ stage1_result_path=stage1_result_path,
294
+ stage2_index=stage2_index,
295
+ seed_2=seed_2,
296
+ guidance_scale_2=guidance_scale_2,
297
+ custom_timesteps_2=custom_timesteps_2,
298
+ num_inference_steps_2=num_inference_steps_2,
299
+ disable_watermark=True)
300
+ temp_image = out_image.copy()
301
+ self.super_res_1_pipe.watermarker.apply_watermark(
302
+ [temp_image], self.super_res_1_pipe.unet.config.sample_size)
303
+ yield temp_image
304
+ yield self.run_stage3(image=out_image,
305
+ prompt=prompt,
306
+ negative_prompt=negative_prompt,
307
+ seed_3=seed_3,
308
+ guidance_scale_3=guidance_scale_3,
309
+ num_inference_steps_3=num_inference_steps_3)
requirements.txt ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ accelerate==0.18.0
2
+ beautifulsoup4==4.12.2
3
+ bitsandbytes==0.38.1
4
+ diffusers==0.16.0
5
+ ftfy==6.1.1
6
+ gradio==3.27.0
7
+ huggingface_hub==0.14.1
8
+ numpy==1.24.3
9
+ Pillow==9.5.0
10
+ safetensors==0.3.0
11
+ sentencepiece==0.1.98
12
+ tokenizers==0.13.3
13
+ torch==2.0.0
14
+ torchvision==0.15.1
15
+ tqdm==4.65.0
16
+ transformers==4.28.1
settings.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import numpy as np
4
+
5
+ HF_TOKEN = os.getenv('HF_TOKEN')
6
+ UPLOAD_REPO_ID = os.getenv('UPLOAD_REPO_ID')
7
+ UPLOAD_RESULT_IMAGE = os.getenv('UPLOAD_RESULT_IMAGE') == '1'
8
+
9
+ # UI options
10
+ SHOW_DUPLICATE_BUTTON = os.getenv('SHOW_DUPLICATE_BUTTON', '0') == '1'
11
+ SHOW_DEVICE_WARNING = os.getenv('SHOW_DEVICE_WARNING', '1') == '1'
12
+ SHOW_ADVANCED_OPTIONS = os.getenv('SHOW_ADVANCED_OPTIONS', '1') == '1'
13
+ SHOW_UPSCALE_TO_256_BUTTON = os.getenv('SHOW_UPSCALE_TO_256_BUTTON',
14
+ '0') == '1'
15
+ SHOW_NUM_IMAGES = os.getenv('SHOW_NUM_IMAGES_OPTION', '1') == '1'
16
+ SHOW_CUSTOM_TIMESTEPS_1 = os.getenv('SHOW_CUSTOM_TIMESTEPS_1', '1') == '1'
17
+ SHOW_CUSTOM_TIMESTEPS_2 = os.getenv('SHOW_CUSTOM_TIMESTEPS_2', '1') == '1'
18
+ SHOW_NUM_STEPS_1 = os.getenv('SHOW_NUM_STEPS_1', '0') == '1'
19
+ SHOW_NUM_STEPS_2 = os.getenv('SHOW_NUM_STEPS_2', '0') == '1'
20
+ SHOW_NUM_STEPS_3 = os.getenv('SHOW_NUM_STEPS_3', '1') == '1'
21
+ GALLERY_COLUMN_NUM = int(os.getenv('GALLERY_COLUMN_NUM', '4'))
22
+
23
+ # Parameters
24
+ MAX_QUEUE_SIZE = int(os.getenv('MAX_QUEUE_SIZE', '10'))
25
+ MAX_SEED = np.iinfo(np.int32).max
26
+ MAX_NUM_IMAGES = int(os.getenv('MAX_NUM_IMAGES', '4'))
27
+ DEFAULT_NUM_IMAGES = min(MAX_NUM_IMAGES,
28
+ int(os.getenv('DEFAULT_NUM_IMAGES', '4')))
29
+ MAX_NUM_STEPS = int(os.getenv('MAX_NUM_STEPS', '200'))
30
+ DEFAULT_CUSTOM_TIMESTEPS_1 = os.getenv('DEFAULT_CUSTOM_TIMESTEPS_1',
31
+ 'smart100')
32
+ DEFAULT_CUSTOM_TIMESTEPS_2 = os.getenv('DEFAULT_CUSTOM_TIMESTEPS_2', 'smart50')
33
+ DEFAULT_NUM_STEPS_3 = int(os.getenv('DEFAULT_NUM_STEPS_3', '40'))
34
+
35
+ # Model options
36
+ DISABLE_AUTOMATIC_CPU_OFFLOAD = os.getenv(
37
+ 'DISABLE_AUTOMATIC_CPU_OFFLOAD') == '1'
38
+ DISABLE_SD_X4_UPSCALER = os.getenv('DISABLE_SD_X4_UPSCALER') == '1'
39
+
40
+ # Other options
41
+ RUN_GARBAGE_COLLECTION = os.getenv('RUN_GARBAGE_COLLECTION', '1') == '1'
42
+ DEBUG = os.getenv('DEBUG') == '1'
43
+
44
+ # Default options for the public demo
45
+ if os.getenv('IS_PUBLIC_DEMO') == '1':
46
+ # UI
47
+ SHOW_DUPLICATE_BUTTON = True
48
+ SHOW_NUM_STEPS_3 = False
49
+ SHOW_CUSTOM_TIMESTEPS_1 = False
50
+ SHOW_CUSTOM_TIMESTEPS_2 = False
51
+ SHOW_NUM_IMAGES = False
52
+ # parameters
53
+ DEFAULT_CUSTOM_TIMESTEPS_1 = 'smart50'
54
+ # model
55
+ DISABLE_AUTOMATIC_CPU_OFFLOAD = True
56
+ RUN_GARBAGE_COLLECTION = False
style.css ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /*
2
+ This CSS file is modified from:
3
+ https://huggingface.co/spaces/stabilityai/stable-diffusion/blob/2794a3c3ba66115c307075098e713f572b08bf80/app.py
4
+ */
5
+
6
+ h1 {
7
+ text-align: center;
8
+ }
9
+
10
+ .gradio-container {
11
+ font-family: 'IBM Plex Sans', sans-serif;
12
+ }
13
+