Duplicate from DeepFloyd/IF

Co-authored-by: hysts <[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
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,162 @@
+gradio_cached_examples/
+
+# 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/

.pre-commit-config.yaml

@@ -0,0 +1,38 @@
+exclude: diffusers-private
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
+      files: ^diffusers-private
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

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

README.md

@@ -0,0 +1,15 @@
+---
+title: IF
+emoji: 🔥
+colorFrom: pink
+colorTo: red
+sdk: gradio
+sdk_version: 3.27.0
+python_version: 3.10.11
+app_file: app.py
+pinned: false
+license: other
+duplicated_from: DeepFloyd/IF
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,664 @@
+#!/usr/bin/env python
+
+import datetime
+import hashlib
+import json
+import os
+import random
+import tempfile
+
+import gradio as gr
+import torch
+from huggingface_hub import HfApi
+
+# isort: off
+from model import Model
+from settings import (
+    DEBUG,
+    DEFAULT_CUSTOM_TIMESTEPS_1,
+    DEFAULT_CUSTOM_TIMESTEPS_2,
+    DEFAULT_NUM_IMAGES,
+    DEFAULT_NUM_STEPS_3,
+    DISABLE_SD_X4_UPSCALER,
+    GALLERY_COLUMN_NUM,
+    HF_TOKEN,
+    MAX_NUM_IMAGES,
+    MAX_NUM_STEPS,
+    MAX_QUEUE_SIZE,
+    MAX_SEED,
+    SHOW_ADVANCED_OPTIONS,
+    SHOW_CUSTOM_TIMESTEPS_1,
+    SHOW_CUSTOM_TIMESTEPS_2,
+    SHOW_DEVICE_WARNING,
+    SHOW_DUPLICATE_BUTTON,
+    SHOW_NUM_IMAGES,
+    SHOW_NUM_STEPS_1,
+    SHOW_NUM_STEPS_2,
+    SHOW_NUM_STEPS_3,
+    SHOW_UPSCALE_TO_256_BUTTON,
+    UPLOAD_REPO_ID,
+    UPLOAD_RESULT_IMAGE,
+)
+# isort: on
+
+TITLE = '# [DeepFloyd IF](https://github.com/deep-floyd/IF)'
+DESCRIPTION = 'The DeepFloyd IF model has been initially released as a non-commercial research-only model. Please make sure you read and abide to the [LICENSE](https://huggingface.co/spaces/DeepFloyd/deepfloyd-if-license) before using it.'
+DISCLAIMER = 'In this demo, the DeepFloyd team may collect prompts, and user preferences (which of the images the user chose to upscale) for improving future models'
+FOOTER = """<div class="footer">
+                    <p>Model by <a href="https://huggingface.co/DeepFloyd" style="text-decoration: underline;" target="_blank">DeepFloyd</a> supported by <a href="https://huggingface.co/stabilityai" style="text-decoration: underline;" target="_blank">Stability AI</a>
+                    </p>
+            </div>
+            <div class="acknowledgments">
+                    <p><h4>LICENSE</h4>
+The model is licensed with a bespoke non-commercial research-only license <a href="https://huggingface.co/spaces/DeepFloyd/deepfloyd-if-license" style="text-decoration: underline;" target="_blank">DeepFloyd IF Research License Agreement</a> license. The license forbids you from sharing any content for commercial use, or that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. For the full list of restrictions please <a href="https://huggingface.co/spaces/DeepFloyd/deepfloyd-if-license" style="text-decoration: underline;" target="_blank">read the license</a></p>
+                    <p><h4>Biases and content acknowledgment</h4>
+Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exacerbates societal biases, as well as realistic faces, explicit content and violence. The model was trained on a subset of the <a href="https://laion.ai/blog/laion-5b/" style="text-decoration: underline;" target="_blank">LAION-5B dataset</a> and is meant for research purposes. You can read more in the <a href="https://huggingface.co/DeepFloyd/IF-I-IF-v1.0" style="text-decoration: underline;" target="_blank">model card</a></p>
+            </div>
+        """
+if SHOW_DUPLICATE_BUTTON:
+    SPACE_ID = os.getenv('SPACE_ID')
+    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>'
+
+if SHOW_DEVICE_WARNING and not torch.cuda.is_available():
+    DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'
+
+model = Model()
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+def get_stage2_index(evt: gr.SelectData) -> int:
+    return evt.index
+
+
+def check_if_stage2_selected(index: int) -> None:
+    if index == -1:
+        raise gr.Error(
+            'You need to select the image you would like to upscale from the Stage 1 results by clicking.'
+        )
+
+
+hf_api = HfApi(token=HF_TOKEN)
+if UPLOAD_REPO_ID:
+    hf_api.create_repo(repo_id=UPLOAD_REPO_ID,
+                       private=True,
+                       repo_type='dataset',
+                       exist_ok=True)
+
+
+def get_param_file_hash_name(param_filepath: str) -> str:
+    if not UPLOAD_REPO_ID:
+        return ''
+    with open(param_filepath, 'rb') as f:
+        md5 = hashlib.md5(f.read()).hexdigest()
+    utcnow = datetime.datetime.utcnow().strftime('%Y-%m-%d-%H-%M-%S-%f')
+    return f'{utcnow}-{md5}'
+
+
+def upload_stage1_result(stage1_param_path: str, stage1_result_path: str,
+                         save_name: str) -> None:
+    if not UPLOAD_REPO_ID:
+        return
+    try:
+        hf_api.upload_file(path_or_fileobj=stage1_param_path,
+                           path_in_repo=f'stage1_params/{save_name}.json',
+                           repo_id=UPLOAD_REPO_ID,
+                           repo_type='dataset')
+        hf_api.upload_file(path_or_fileobj=stage1_result_path,
+                           path_in_repo=f'stage1_results/{save_name}.pth',
+                           repo_id=UPLOAD_REPO_ID,
+                           repo_type='dataset')
+    except Exception as e:
+        print(e)
+
+
+def upload_stage2_info(stage1_param_file_hash_name: str,
+                       stage2_output_path: str,
+                       selected_index_for_upscale: int, seed_2: int,
+                       guidance_scale_2: float, custom_timesteps_2: str,
+                       num_inference_steps_2: int) -> None:
+    if not UPLOAD_REPO_ID:
+        return
+    if not stage1_param_file_hash_name:
+        raise ValueError
+
+    stage2_params = {
+        'stage1_param_file_hash_name': stage1_param_file_hash_name,
+        'selected_index_for_upscale': selected_index_for_upscale,
+        'seed_2': seed_2,
+        'guidance_scale_2': guidance_scale_2,
+        'custom_timesteps_2': custom_timesteps_2,
+        'num_inference_steps_2': num_inference_steps_2,
+    }
+    with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
+        param_file.write(json.dumps(stage2_params))
+    stage2_param_file_hash_name = get_param_file_hash_name(param_file.name)
+    save_name = f'{stage1_param_file_hash_name}_{stage2_param_file_hash_name}'
+
+    try:
+        hf_api.upload_file(path_or_fileobj=param_file.name,
+                           path_in_repo=f'stage2_params/{save_name}.json',
+                           repo_id=UPLOAD_REPO_ID,
+                           repo_type='dataset')
+        if UPLOAD_RESULT_IMAGE:
+            hf_api.upload_file(path_or_fileobj=stage2_output_path,
+                               path_in_repo=f'stage2_results/{save_name}.png',
+                               repo_id=UPLOAD_REPO_ID,
+                               repo_type='dataset')
+    except Exception as e:
+        print(e)
+
+
+def upload_stage2_3_info(stage1_param_file_hash_name: str,
+                         stage2_3_output_path: str,
+                         selected_index_for_upscale: int, seed_2: int,
+                         guidance_scale_2: float, custom_timesteps_2: str,
+                         num_inference_steps_2: int, prompt: str,
+                         negative_prompt: str, seed_3: int,
+                         guidance_scale_3: float,
+                         num_inference_steps_3: int) -> None:
+    if not UPLOAD_REPO_ID:
+        return
+    if not stage1_param_file_hash_name:
+        raise ValueError
+
+    stage2_3_params = {
+        'stage1_param_file_hash_name': stage1_param_file_hash_name,
+        'selected_index_for_upscale': selected_index_for_upscale,
+        'seed_2': seed_2,
+        'guidance_scale_2': guidance_scale_2,
+        'custom_timesteps_2': custom_timesteps_2,
+        'num_inference_steps_2': num_inference_steps_2,
+        'prompt': prompt,
+        'negative_prompt': negative_prompt,
+        'seed_3': seed_3,
+        'guidance_scale_3': guidance_scale_3,
+        'num_inference_steps_3': num_inference_steps_3,
+    }
+    with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
+        param_file.write(json.dumps(stage2_3_params))
+    stage2_3_param_file_hash_name = get_param_file_hash_name(param_file.name)
+    save_name = f'{stage1_param_file_hash_name}_{stage2_3_param_file_hash_name}'
+
+    try:
+        hf_api.upload_file(path_or_fileobj=param_file.name,
+                           path_in_repo=f'stage2_3_params/{save_name}.json',
+                           repo_id=UPLOAD_REPO_ID,
+                           repo_type='dataset')
+        if UPLOAD_RESULT_IMAGE:
+            hf_api.upload_file(
+                path_or_fileobj=stage2_3_output_path,
+                path_in_repo=f'stage2_3_results/{save_name}.png',
+                repo_id=UPLOAD_REPO_ID,
+                repo_type='dataset')
+    except Exception as e:
+        print(e)
+
+
+def update_upscale_button(selected_index: int) -> tuple[dict, dict]:
+    if selected_index == -1:
+        return gr.update(interactive=False), gr.update(interactive=False)
+    else:
+        return gr.update(interactive=True), gr.update(interactive=True)
+
+
+def _update_result_view(show_gallery: bool) -> tuple[dict, dict]:
+    return gr.update(visible=show_gallery), gr.update(visible=not show_gallery)
+
+
+def show_gallery_view() -> tuple[dict, dict]:
+    return _update_result_view(True)
+
+
+def show_upscaled_view() -> tuple[dict, dict]:
+    return _update_result_view(False)
+
+
+examples = [
+    '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',
+    'paper quilling, extremely detailed, paper quilling of a nordic mountain landscape, 8k rendering',
+    'letters made of candy on a plate that says "diet"',
+    '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"',
+    'ultra close-up color photo portrait of rainbow owl with deer horns in the woods',
+    'a cloth embroidered with the text "laion" and an embroidered cute baby lion face',
+    'product image of a crochet Cthulhu the great old one emerging from a spacetime wormhole made of wool',
+    'a little green budgie parrot driving small red toy car in new york street, photo',
+    'origami dancer in white paper, 3d render, ultra-detailed, on white background, studio shot.',
+    'glowing mushrooms in a natural environment with smoke in the frame',
+    'a subway train\'s digital sign saying "open source", vsco preset, 35mm photo, film grain, in a dim subway station',
+    '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',
+    'a face of a woman made completely out of foliage, twigs, leaves and flowers, side view'
+]
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(TITLE)
+    gr.Markdown(DESCRIPTION)
+    with gr.Box():
+        with gr.Row(elem_id='prompt-container').style(equal_height=True):
+            with gr.Column():
+                prompt = gr.Text(
+                    label='Prompt',
+                    show_label=False,
+                    max_lines=1,
+                    placeholder='Enter your prompt',
+                    elem_id='prompt-text-input',
+                ).style(container=False)
+                negative_prompt = gr.Text(
+                    label='Negative prompt',
+                    show_label=False,
+                    max_lines=1,
+                    placeholder='Enter a negative prompt',
+                    elem_id='negative-prompt-text-input',
+                ).style(container=False)
+            generate_button = gr.Button('Generate').style(full_width=False)
+
+        with gr.Column() as gallery_view:
+            gallery = gr.Gallery(label='Stage 1 results',
+                                 show_label=False,
+                                 elem_id='gallery').style(
+                                     columns=GALLERY_COLUMN_NUM,
+                                     object_fit='contain')
+            gr.Markdown('Pick your favorite generation to upscale.')
+            with gr.Row():
+                upscale_to_256_button = gr.Button(
+                    'Upscale to 256px',
+                    visible=SHOW_UPSCALE_TO_256_BUTTON
+                    or DISABLE_SD_X4_UPSCALER,
+                    interactive=False)
+                upscale_button = gr.Button('Upscale',
+                                           interactive=False,
+                                           visible=not DISABLE_SD_X4_UPSCALER)
+        with gr.Column(visible=False) as upscale_view:
+            result = gr.Image(label='Result',
+                              show_label=False,
+                              type='filepath',
+                              interactive=False,
+                              elem_id='upscaled-image').style(height=640)
+            back_to_selection_button = gr.Button('Back to selection')
+
+        with gr.Accordion('Advanced options',
+                          open=False,
+                          visible=SHOW_ADVANCED_OPTIONS):
+            with gr.Tabs():
+                with gr.Tab(label='Generation'):
+                    seed_1 = gr.Slider(label='Seed',
+                                       minimum=0,
+                                       maximum=MAX_SEED,
+                                       step=1,
+                                       value=0)
+                    randomize_seed_1 = gr.Checkbox(label='Randomize seed',
+                                                   value=True)
+                    guidance_scale_1 = gr.Slider(label='Guidance scale',
+                                                 minimum=1,
+                                                 maximum=20,
+                                                 step=0.1,
+                                                 value=7.0)
+                    custom_timesteps_1 = gr.Dropdown(
+                        label='Custom timesteps 1',
+                        choices=[
+                            'none',
+                            'fast27',
+                            'smart27',
+                            'smart50',
+                            'smart100',
+                            'smart185',
+                        ],
+                        value=DEFAULT_CUSTOM_TIMESTEPS_1,
+                        visible=SHOW_CUSTOM_TIMESTEPS_1)
+                    num_inference_steps_1 = gr.Slider(
+                        label='Number of inference steps',
+                        minimum=1,
+                        maximum=MAX_NUM_STEPS,
+                        step=1,
+                        value=100,
+                        visible=SHOW_NUM_STEPS_1)
+                    num_images = gr.Slider(label='Number of images',
+                                           minimum=1,
+                                           maximum=MAX_NUM_IMAGES,
+                                           step=1,
+                                           value=DEFAULT_NUM_IMAGES,
+                                           visible=SHOW_NUM_IMAGES)
+                with gr.Tab(label='Super-resolution 1'):
+                    seed_2 = gr.Slider(label='Seed',
+                                       minimum=0,
+                                       maximum=MAX_SEED,
+                                       step=1,
+                                       value=0)
+                    randomize_seed_2 = gr.Checkbox(label='Randomize seed',
+                                                   value=True)
+                    guidance_scale_2 = gr.Slider(label='Guidance scale',
+                                                 minimum=1,
+                                                 maximum=20,
+                                                 step=0.1,
+                                                 value=4.0)
+                    custom_timesteps_2 = gr.Dropdown(
+                        label='Custom timesteps 2',
+                        choices=[
+                            'none',
+                            'fast27',
+                            'smart27',
+                            'smart50',
+                            'smart100',
+                            'smart185',
+                        ],
+                        value=DEFAULT_CUSTOM_TIMESTEPS_2,
+                        visible=SHOW_CUSTOM_TIMESTEPS_2)
+                    num_inference_steps_2 = gr.Slider(
+                        label='Number of inference steps',
+                        minimum=1,
+                        maximum=MAX_NUM_STEPS,
+                        step=1,
+                        value=50,
+                        visible=SHOW_NUM_STEPS_2)
+                with gr.Tab(label='Super-resolution 2'):
+                    seed_3 = gr.Slider(label='Seed',
+                                       minimum=0,
+                                       maximum=MAX_SEED,
+                                       step=1,
+                                       value=0)
+                    randomize_seed_3 = gr.Checkbox(label='Randomize seed',
+                                                   value=True)
+                    guidance_scale_3 = gr.Slider(label='Guidance scale',
+                                                 minimum=1,
+                                                 maximum=20,
+                                                 step=0.1,
+                                                 value=9.0)
+                    num_inference_steps_3 = gr.Slider(
+                        label='Number of inference steps',
+                        minimum=1,
+                        maximum=MAX_NUM_STEPS,
+                        step=1,
+                        value=DEFAULT_NUM_STEPS_3,
+                        visible=SHOW_NUM_STEPS_3)
+
+    gr.Examples(examples=examples, inputs=prompt, examples_per_page=4)
+
+    with gr.Box(visible=DEBUG):
+        with gr.Row():
+            with gr.Accordion(label='Hidden params'):
+                stage1_param_path = gr.Text(label='Stage 1 param path')
+                stage1_result_path = gr.Text(label='Stage 1 result path')
+                stage1_param_file_hash_name = gr.Text(
+                    label='Stage 1 param file hash name')
+                selected_index_for_stage2 = gr.Number(
+                    label='Selected index for Stage 2', value=-1, precision=0)
+    gr.Markdown(DISCLAIMER)
+    gr.HTML(FOOTER)
+    stage1_inputs = [
+        prompt,
+        negative_prompt,
+        seed_1,
+        num_images,
+        guidance_scale_1,
+        custom_timesteps_1,
+        num_inference_steps_1,
+    ]
+    stage1_outputs = [
+        gallery,
+        stage1_param_path,
+        stage1_result_path,
+    ]
+
+    prompt.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed_1, randomize_seed_1],
+        outputs=seed_1,
+        queue=False,
+    ).then(
+        fn=lambda: -1,
+        outputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=show_gallery_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=update_upscale_button,
+        inputs=selected_index_for_stage2,
+        outputs=[
+            upscale_button,
+            upscale_to_256_button,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage1,
+        inputs=stage1_inputs,
+        outputs=stage1_outputs,
+    ).success(
+        fn=get_param_file_hash_name,
+        inputs=stage1_param_path,
+        outputs=stage1_param_file_hash_name,
+        queue=False,
+    ).then(
+        fn=upload_stage1_result,
+        inputs=[
+            stage1_param_path,
+            stage1_result_path,
+            stage1_param_file_hash_name,
+        ],
+        queue=False,
+    )
+
+    negative_prompt.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed_1, randomize_seed_1],
+        outputs=seed_1,
+        queue=False,
+    ).then(
+        fn=lambda: -1,
+        outputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=show_gallery_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=update_upscale_button,
+        inputs=selected_index_for_stage2,
+        outputs=[
+            upscale_button,
+            upscale_to_256_button,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage1,
+        inputs=stage1_inputs,
+        outputs=stage1_outputs,
+    ).success(
+        fn=get_param_file_hash_name,
+        inputs=stage1_param_path,
+        outputs=stage1_param_file_hash_name,
+        queue=False,
+    ).then(
+        fn=upload_stage1_result,
+        inputs=[
+            stage1_param_path,
+            stage1_result_path,
+            stage1_param_file_hash_name,
+        ],
+        queue=False,
+    )
+
+    generate_button.click(
+        fn=randomize_seed_fn,
+        inputs=[seed_1, randomize_seed_1],
+        outputs=seed_1,
+        queue=False,
+    ).then(
+        fn=lambda: -1,
+        outputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=show_gallery_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=update_upscale_button,
+        inputs=selected_index_for_stage2,
+        outputs=[
+            upscale_button,
+            upscale_to_256_button,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage1,
+        inputs=stage1_inputs,
+        outputs=stage1_outputs,
+        api_name='generate64',
+    ).success(
+        fn=get_param_file_hash_name,
+        inputs=stage1_param_path,
+        outputs=stage1_param_file_hash_name,
+        queue=False,
+    ).then(
+        fn=upload_stage1_result,
+        inputs=[
+            stage1_param_path,
+            stage1_result_path,
+            stage1_param_file_hash_name,
+        ],
+        queue=False,
+    )
+
+    gallery.select(
+        fn=get_stage2_index,
+        outputs=selected_index_for_stage2,
+        queue=False,
+    )
+
+    selected_index_for_stage2.change(
+        fn=update_upscale_button,
+        inputs=selected_index_for_stage2,
+        outputs=[
+            upscale_button,
+            upscale_to_256_button,
+        ],
+        queue=False,
+    )
+
+    stage2_inputs = [
+        stage1_result_path,
+        selected_index_for_stage2,
+        seed_2,
+        guidance_scale_2,
+        custom_timesteps_2,
+        num_inference_steps_2,
+    ]
+
+    upscale_to_256_button.click(
+        fn=check_if_stage2_selected,
+        inputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=randomize_seed_fn,
+        inputs=[seed_2, randomize_seed_2],
+        outputs=seed_2,
+        queue=False,
+    ).then(
+        fn=show_upscaled_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage2,
+        inputs=stage2_inputs,
+        outputs=result,
+        api_name='upscale256',
+    ).success(
+        fn=upload_stage2_info,
+        inputs=[
+            stage1_param_file_hash_name,
+            result,
+            selected_index_for_stage2,
+            seed_2,
+            guidance_scale_2,
+            custom_timesteps_2,
+            num_inference_steps_2,
+        ],
+        queue=False,
+    )
+
+    stage2_3_inputs = [
+        stage1_result_path,
+        selected_index_for_stage2,
+        seed_2,
+        guidance_scale_2,
+        custom_timesteps_2,
+        num_inference_steps_2,
+        prompt,
+        negative_prompt,
+        seed_3,
+        guidance_scale_3,
+        num_inference_steps_3,
+    ]
+
+    upscale_button.click(
+        fn=check_if_stage2_selected,
+        inputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=randomize_seed_fn,
+        inputs=[seed_2, randomize_seed_2],
+        outputs=seed_2,
+        queue=False,
+    ).then(
+        fn=randomize_seed_fn,
+        inputs=[seed_3, randomize_seed_3],
+        outputs=seed_3,
+        queue=False,
+    ).then(
+        fn=show_upscaled_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage2_3,
+        inputs=stage2_3_inputs,
+        outputs=result,
+        api_name='upscale1024',
+    ).success(
+        fn=upload_stage2_3_info,
+        inputs=[
+            stage1_param_file_hash_name,
+            result,
+            selected_index_for_stage2,
+            seed_2,
+            guidance_scale_2,
+            custom_timesteps_2,
+            num_inference_steps_2,
+            prompt,
+            negative_prompt,
+            seed_3,
+            guidance_scale_3,
+            num_inference_steps_3,
+        ],
+        queue=False,
+    )
+
+    back_to_selection_button.click(
+        fn=show_gallery_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    )
+
+demo.queue(api_open=False, max_size=MAX_QUEUE_SIZE).launch(debug=DEBUG)

model.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import gc
+import json
+import tempfile
+from typing import Generator
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers import DiffusionPipeline, StableDiffusionUpscalePipeline
+from diffusers.pipelines.deepfloyd_if import (fast27_timesteps,
+                                              smart27_timesteps,
+                                              smart50_timesteps,
+                                              smart100_timesteps,
+                                              smart185_timesteps)
+
+from settings import (DISABLE_AUTOMATIC_CPU_OFFLOAD, DISABLE_SD_X4_UPSCALER,
+                      HF_TOKEN, MAX_NUM_IMAGES, MAX_NUM_STEPS, MAX_SEED,
+                      RUN_GARBAGE_COLLECTION)
+
+
+class Model:
+    def __init__(self):
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.pipe = None
+        self.super_res_1_pipe = None
+        self.super_res_2_pipe = None
+        self.watermark_image = None
+
+        if torch.cuda.is_available():
+            self.load_weights()
+            self.watermark_image = PIL.Image.fromarray(
+                self.pipe.watermarker.watermark_image.to(
+                    torch.uint8).cpu().numpy(),
+                mode='RGBA')
+
+    def load_weights(self) -> None:
+        self.pipe = DiffusionPipeline.from_pretrained(
+            'DeepFloyd/IF-I-XL-v1.0',
+            torch_dtype=torch.float16,
+            variant='fp16',
+            use_safetensors=True,
+            use_auth_token=HF_TOKEN)
+        self.super_res_1_pipe = DiffusionPipeline.from_pretrained(
+            'DeepFloyd/IF-II-L-v1.0',
+            text_encoder=None,
+            torch_dtype=torch.float16,
+            variant='fp16',
+            use_safetensors=True,
+            use_auth_token=HF_TOKEN)
+
+        if not DISABLE_SD_X4_UPSCALER:
+            self.super_res_2_pipe = StableDiffusionUpscalePipeline.from_pretrained(
+                'stabilityai/stable-diffusion-x4-upscaler',
+                torch_dtype=torch.float16)
+
+        if DISABLE_AUTOMATIC_CPU_OFFLOAD:
+            self.pipe.to(self.device)
+            self.super_res_1_pipe.to(self.device)
+            if not DISABLE_SD_X4_UPSCALER:
+                self.super_res_2_pipe.to(self.device)
+        else:
+            self.pipe.enable_model_cpu_offload()
+            self.super_res_1_pipe.enable_model_cpu_offload()
+            if not DISABLE_SD_X4_UPSCALER:
+                self.super_res_2_pipe.enable_model_cpu_offload()
+
+    def apply_watermark_to_sd_x4_upscaler_results(
+            self, images: list[PIL.Image.Image]) -> None:
+        w, h = images[0].size
+
+        stability_x4_upscaler_sample_size = 128
+
+        coef = min(h / stability_x4_upscaler_sample_size,
+                   w / stability_x4_upscaler_sample_size)
+        img_h, img_w = (int(h / coef), int(w / coef)) if coef < 1 else (h, w)
+
+        S1, S2 = 1024**2, img_w * img_h
+        K = (S2 / S1)**0.5
+        watermark_size = int(K * 62)
+        watermark_x = img_w - int(14 * K)
+        watermark_y = img_h - int(14 * K)
+
+        watermark_image = self.watermark_image.copy().resize(
+            (watermark_size, watermark_size),
+            PIL.Image.Resampling.BICUBIC,
+            reducing_gap=None)
+
+        for image in images:
+            image.paste(watermark_image,
+                        box=(
+                            watermark_x - watermark_size,
+                            watermark_y - watermark_size,
+                            watermark_x,
+                            watermark_y,
+                        ),
+                        mask=watermark_image.split()[-1])
+
+    @staticmethod
+    def to_pil_images(images: torch.Tensor) -> list[PIL.Image.Image]:
+        images = (images / 2 + 0.5).clamp(0, 1)
+        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+        images = np.round(images * 255).astype(np.uint8)
+        return [PIL.Image.fromarray(image) for image in images]
+
+    @staticmethod
+    def check_seed(seed: int) -> None:
+        if not 0 <= seed <= MAX_SEED:
+            raise ValueError
+
+    @staticmethod
+    def check_num_images(num_images: int) -> None:
+        if not 1 <= num_images <= MAX_NUM_IMAGES:
+            raise ValueError
+
+    @staticmethod
+    def check_num_inference_steps(num_steps: int) -> None:
+        if not 1 <= num_steps <= MAX_NUM_STEPS:
+            raise ValueError
+
+    @staticmethod
+    def get_custom_timesteps(name: str) -> list[int] | None:
+        if name == 'none':
+            timesteps = None
+        elif name == 'fast27':
+            timesteps = fast27_timesteps
+        elif name == 'smart27':
+            timesteps = smart27_timesteps
+        elif name == 'smart50':
+            timesteps = smart50_timesteps
+        elif name == 'smart100':
+            timesteps = smart100_timesteps
+        elif name == 'smart185':
+            timesteps = smart185_timesteps
+        else:
+            raise ValueError
+        return timesteps
+
+    @staticmethod
+    def run_garbage_collection():
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def run_stage1(
+        self,
+        prompt: str,
+        negative_prompt: str = '',
+        seed: int = 0,
+        num_images: int = 1,
+        guidance_scale_1: float = 7.0,
+        custom_timesteps_1: str = 'smart100',
+        num_inference_steps_1: int = 100,
+    ) -> tuple[list[PIL.Image.Image], str, str]:
+        self.check_seed(seed)
+        self.check_num_images(num_images)
+        self.check_num_inference_steps(num_inference_steps_1)
+
+        if RUN_GARBAGE_COLLECTION:
+            self.run_garbage_collection()
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        prompt_embeds, negative_embeds = self.pipe.encode_prompt(
+            prompt=prompt, negative_prompt=negative_prompt)
+
+        timesteps = self.get_custom_timesteps(custom_timesteps_1)
+
+        images = self.pipe(prompt_embeds=prompt_embeds,
+                           negative_prompt_embeds=negative_embeds,
+                           num_images_per_prompt=num_images,
+                           guidance_scale=guidance_scale_1,
+                           timesteps=timesteps,
+                           num_inference_steps=num_inference_steps_1,
+                           generator=generator,
+                           output_type='pt').images
+        pil_images = self.to_pil_images(images)
+        self.pipe.watermarker.apply_watermark(
+            pil_images, self.pipe.unet.config.sample_size)
+
+        stage1_params = {
+            'prompt': prompt,
+            'negative_prompt': negative_prompt,
+            'seed': seed,
+            'num_images': num_images,
+            'guidance_scale_1': guidance_scale_1,
+            'custom_timesteps_1': custom_timesteps_1,
+            'num_inference_steps_1': num_inference_steps_1,
+        }
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file:
+            param_file.write(json.dumps(stage1_params))
+        stage1_result = {
+            'prompt_embeds': prompt_embeds,
+            'negative_embeds': negative_embeds,
+            'images': images,
+            'pil_images': pil_images,
+        }
+        with tempfile.NamedTemporaryFile(delete=False) as result_file:
+            torch.save(stage1_result, result_file.name)
+        return pil_images, param_file.name, result_file.name
+
+    def run_stage2(
+        self,
+        stage1_result_path: str,
+        stage2_index: int,
+        seed_2: int = 0,
+        guidance_scale_2: float = 4.0,
+        custom_timesteps_2: str = 'smart50',
+        num_inference_steps_2: int = 50,
+        disable_watermark: bool = False,
+    ) -> PIL.Image.Image:
+        self.check_seed(seed_2)
+        self.check_num_inference_steps(num_inference_steps_2)
+
+        if RUN_GARBAGE_COLLECTION:
+            self.run_garbage_collection()
+
+        generator = torch.Generator(device=self.device).manual_seed(seed_2)
+
+        stage1_result = torch.load(stage1_result_path)
+        prompt_embeds = stage1_result['prompt_embeds']
+        negative_embeds = stage1_result['negative_embeds']
+        images = stage1_result['images']
+        images = images[[stage2_index]]
+
+        timesteps = self.get_custom_timesteps(custom_timesteps_2)
+
+        out = self.super_res_1_pipe(image=images,
+                                    prompt_embeds=prompt_embeds,
+                                    negative_prompt_embeds=negative_embeds,
+                                    num_images_per_prompt=1,
+                                    guidance_scale=guidance_scale_2,
+                                    timesteps=timesteps,
+                                    num_inference_steps=num_inference_steps_2,
+                                    generator=generator,
+                                    output_type='pt',
+                                    noise_level=250).images
+        pil_images = self.to_pil_images(out)
+
+        if disable_watermark:
+            return pil_images[0]
+
+        self.super_res_1_pipe.watermarker.apply_watermark(
+            pil_images, self.super_res_1_pipe.unet.config.sample_size)
+        return pil_images[0]
+
+    def run_stage3(
+        self,
+        image: PIL.Image.Image,
+        prompt: str = '',
+        negative_prompt: str = '',
+        seed_3: int = 0,
+        guidance_scale_3: float = 9.0,
+        num_inference_steps_3: int = 75,
+    ) -> PIL.Image.Image:
+        self.check_seed(seed_3)
+        self.check_num_inference_steps(num_inference_steps_3)
+
+        if RUN_GARBAGE_COLLECTION:
+            self.run_garbage_collection()
+
+        generator = torch.Generator(device=self.device).manual_seed(seed_3)
+        out = self.super_res_2_pipe(image=image,
+                                    prompt=prompt,
+                                    negative_prompt=negative_prompt,
+                                    num_images_per_prompt=1,
+                                    guidance_scale=guidance_scale_3,
+                                    num_inference_steps=num_inference_steps_3,
+                                    generator=generator,
+                                    noise_level=100).images
+        self.apply_watermark_to_sd_x4_upscaler_results(out)
+        return out[0]
+
+    def run_stage2_3(
+        self,
+        stage1_result_path: str,
+        stage2_index: int,
+        seed_2: int = 0,
+        guidance_scale_2: float = 4.0,
+        custom_timesteps_2: str = 'smart50',
+        num_inference_steps_2: int = 50,
+        prompt: str = '',
+        negative_prompt: str = '',
+        seed_3: int = 0,
+        guidance_scale_3: float = 9.0,
+        num_inference_steps_3: int = 75,
+    ) -> Generator[PIL.Image.Image]:
+        self.check_seed(seed_3)
+        self.check_num_inference_steps(num_inference_steps_3)
+
+        out_image = self.run_stage2(
+            stage1_result_path=stage1_result_path,
+            stage2_index=stage2_index,
+            seed_2=seed_2,
+            guidance_scale_2=guidance_scale_2,
+            custom_timesteps_2=custom_timesteps_2,
+            num_inference_steps_2=num_inference_steps_2,
+            disable_watermark=True)
+        temp_image = out_image.copy()
+        self.super_res_1_pipe.watermarker.apply_watermark(
+            [temp_image], self.super_res_1_pipe.unet.config.sample_size)
+        yield temp_image
+        yield self.run_stage3(image=out_image,
+                              prompt=prompt,
+                              negative_prompt=negative_prompt,
+                              seed_3=seed_3,
+                              guidance_scale_3=guidance_scale_3,
+                              num_inference_steps_3=num_inference_steps_3)

requirements.txt

@@ -0,0 +1,16 @@
+accelerate==0.18.0
+beautifulsoup4==4.12.2
+bitsandbytes==0.38.1
+diffusers==0.16.0
+ftfy==6.1.1
+gradio==3.27.0
+huggingface_hub==0.14.1
+numpy==1.24.3
+Pillow==9.5.0
+safetensors==0.3.0
+sentencepiece==0.1.98
+tokenizers==0.13.3
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.65.0
+transformers==4.28.1

settings.py

@@ -0,0 +1,56 @@
+import os
+
+import numpy as np
+
+HF_TOKEN = os.getenv('HF_TOKEN')
+UPLOAD_REPO_ID = os.getenv('UPLOAD_REPO_ID')
+UPLOAD_RESULT_IMAGE = os.getenv('UPLOAD_RESULT_IMAGE') == '1'
+
+# UI options
+SHOW_DUPLICATE_BUTTON = os.getenv('SHOW_DUPLICATE_BUTTON', '0') == '1'
+SHOW_DEVICE_WARNING = os.getenv('SHOW_DEVICE_WARNING', '1') == '1'
+SHOW_ADVANCED_OPTIONS = os.getenv('SHOW_ADVANCED_OPTIONS', '1') == '1'
+SHOW_UPSCALE_TO_256_BUTTON = os.getenv('SHOW_UPSCALE_TO_256_BUTTON',
+                                       '0') == '1'
+SHOW_NUM_IMAGES = os.getenv('SHOW_NUM_IMAGES_OPTION', '1') == '1'
+SHOW_CUSTOM_TIMESTEPS_1 = os.getenv('SHOW_CUSTOM_TIMESTEPS_1', '1') == '1'
+SHOW_CUSTOM_TIMESTEPS_2 = os.getenv('SHOW_CUSTOM_TIMESTEPS_2', '1') == '1'
+SHOW_NUM_STEPS_1 = os.getenv('SHOW_NUM_STEPS_1', '0') == '1'
+SHOW_NUM_STEPS_2 = os.getenv('SHOW_NUM_STEPS_2', '0') == '1'
+SHOW_NUM_STEPS_3 = os.getenv('SHOW_NUM_STEPS_3', '1') == '1'
+GALLERY_COLUMN_NUM = int(os.getenv('GALLERY_COLUMN_NUM', '4'))
+
+# Parameters
+MAX_QUEUE_SIZE = int(os.getenv('MAX_QUEUE_SIZE', '10'))
+MAX_SEED = np.iinfo(np.int32).max
+MAX_NUM_IMAGES = int(os.getenv('MAX_NUM_IMAGES', '4'))
+DEFAULT_NUM_IMAGES = min(MAX_NUM_IMAGES,
+                         int(os.getenv('DEFAULT_NUM_IMAGES', '4')))
+MAX_NUM_STEPS = int(os.getenv('MAX_NUM_STEPS', '200'))
+DEFAULT_CUSTOM_TIMESTEPS_1 = os.getenv('DEFAULT_CUSTOM_TIMESTEPS_1',
+                                       'smart100')
+DEFAULT_CUSTOM_TIMESTEPS_2 = os.getenv('DEFAULT_CUSTOM_TIMESTEPS_2', 'smart50')
+DEFAULT_NUM_STEPS_3 = int(os.getenv('DEFAULT_NUM_STEPS_3', '40'))
+
+# Model options
+DISABLE_AUTOMATIC_CPU_OFFLOAD = os.getenv(
+    'DISABLE_AUTOMATIC_CPU_OFFLOAD') == '1'
+DISABLE_SD_X4_UPSCALER = os.getenv('DISABLE_SD_X4_UPSCALER') == '1'
+
+# Other options
+RUN_GARBAGE_COLLECTION = os.getenv('RUN_GARBAGE_COLLECTION', '1') == '1'
+DEBUG = os.getenv('DEBUG') == '1'
+
+# Default options for the public demo
+if os.getenv('IS_PUBLIC_DEMO') == '1':
+    # UI
+    SHOW_DUPLICATE_BUTTON = True
+    SHOW_NUM_STEPS_3 = False
+    SHOW_CUSTOM_TIMESTEPS_1 = False
+    SHOW_CUSTOM_TIMESTEPS_2 = False
+    SHOW_NUM_IMAGES = False
+    # parameters
+    DEFAULT_CUSTOM_TIMESTEPS_1 = 'smart50'
+    # model
+    DISABLE_AUTOMATIC_CPU_OFFLOAD = True
+    RUN_GARBAGE_COLLECTION = False

style.css

@@ -0,0 +1,203 @@
+/*
+This CSS file is modified from:
+https://huggingface.co/spaces/stabilityai/stable-diffusion/blob/2794a3c3ba66115c307075098e713f572b08bf80/app.py
+*/
+
+h1 {
+  text-align: center;
+}
+
+.gradio-container {
+  font-family: 'IBM Plex Sans', sans-serif;
+}
+
+.gr-button {
+  color: white;
+  border-color: black;
+  background: black;
+}
+
+input[type='range'] {
+  accent-color: black;
+}
+
+.dark input[type='range'] {
+  accent-color: #dfdfdf;
+}
+
+.container {
+  max-width: 730px;
+  margin: auto;
+  padding-top: 1.5rem;
+}
+
+#gallery {
+  min-height: auto;
+  height: 185px;
+  margin-top: 15px;
+  margin-left: auto;
+  margin-right: auto;
+  border-bottom-right-radius: .5rem !important;
+  border-bottom-left-radius: .5rem !important;
+}
+#gallery .grid-wrap, #gallery .empty{
+  height: 185px;
+  min-height: 185px;
+}
+#gallery .preview{
+  height: 185px;
+  min-height: 185px!important;
+}
+#gallery>div>.h-full {
+  min-height: 20rem;
+}
+
+.details:hover {
+  text-decoration: underline;
+}
+
+.gr-button {
+  white-space: nowrap;
+}
+
+.gr-button:focus {
+  border-color: rgb(147 197 253 / var(--tw-border-opacity));
+  outline: none;
+  box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
+  --tw-border-opacity: 1;
+  --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
+  --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
+  --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
+  --tw-ring-opacity: .5;
+}
+
+#advanced-btn {
+  font-size: .7rem !important;
+  line-height: 19px;
+  margin-top: 12px;
+  margin-bottom: 12px;
+  padding: 2px 8px;
+  border-radius: 14px !important;
+}
+
+#advanced-options {
+  display: none;
+  margin-bottom: 20px;
+}
+
+.footer {
+  margin-bottom: 45px;
+  margin-top: 35px;
+  text-align: center;
+  border-bottom: 1px solid #e5e5e5;
+}
+
+.footer>p {
+  font-size: .8rem;
+  display: inline-block;
+  padding: 0 10px;
+  transform: translateY(10px);
+  background: white;
+}
+
+.dark .footer {
+  border-color: #303030;
+}
+
+.dark .footer>p {
+  background: #0b0f19;
+}
+
+.acknowledgments h4 {
+  margin: 1.25em 0 .25em 0;
+  font-weight: bold;
+  font-size: 115%;
+}
+
+.animate-spin {
+  animation: spin 1s linear infinite;
+}
+
+@keyframes spin {
+  from {
+    transform: rotate(0deg);
+  }
+
+  to {
+    transform: rotate(360deg);
+  }
+}
+
+#share-btn-container {
+  display: flex;
+  padding-left: 0.5rem !important;
+  padding-right: 0.5rem !important;
+  background-color: #000000;
+  justify-content: center;
+  align-items: center;
+  border-radius: 9999px !important;
+  width: 13rem;
+  margin-top: 10px;
+  margin-left: auto;
+}
+
+#share-btn {
+  all: initial;
+  color: #ffffff;
+  font-weight: 600;
+  cursor: pointer;
+  font-family: 'IBM Plex Sans', sans-serif;
+  margin-left: 0.5rem !important;
+  padding-top: 0.25rem !important;
+  padding-bottom: 0.25rem !important;
+  right: 0;
+}
+
+#share-btn * {
+  all: unset;
+}
+
+#share-btn-container div:nth-child(-n+2) {
+  width: auto !important;
+  min-height: 0px !important;
+}
+
+#share-btn-container .wrap {
+  display: none !important;
+}
+
+.gr-form {
+  flex: 1 1 50%;
+  border-top-right-radius: 0;
+  border-bottom-right-radius: 0;
+}
+
+#prompt-container {
+  gap: 0;
+}
+
+#prompt-text-input,
+#negative-prompt-text-input {
+  padding: .45rem 0.625rem
+}
+
+#component-16 {
+  border-top-width: 1px !important;
+  margin-top: 1em
+}
+
+.image_duplication {
+  position: absolute;
+  width: 100px;
+  left: 50px
+}
+
+#component-0 {
+  max-width: 730px;
+  margin: auto;
+  padding-top: 1.5rem;
+}
+
+#upscaled-image img {
+  object-fit: scale-down;
+}