From 5a48c2a6e7c52ec044d2eb8219ce2f3ef12be458 Mon Sep 17 00:00:00 2001 From: ddPn08 Date: Sat, 29 Apr 2023 21:47:53 +0900 Subject: [PATCH] use gradio --- .gitignore | 200 ++++++++++++++++-- api/base.py | 6 - api/models/diffusion.py | 18 +- api/models/plugin.py | 1 - api/models/sockets.py | 10 - api/models/tensorrt.py | 2 +- bin/.gitignore | 2 + dev.py | 5 + launch.py | 131 +++++++----- lib/diffusers/scheduler.py | 14 ++ lib/tensorrt/models.py | 29 ++- lib/tensorrt/pipeline.py | 17 +- lib/tensorrt/pipeline_stable_diffusion.py | 2 +- .../pipeline_stable_diffusion_img2img.py | 2 +- lib/tensorrt/utilities.py | 19 +- modules/acceleration/tensorrt/engine.py | 80 +------ modules/acceleration/tensorrt/text_encoder.py | 1 - modules/app.py | 35 --- modules/components/header.py | 140 ++++++++++++ .../components/image_generation_options.py | 94 ++++++++ modules/config.py | 26 ++- modules/diffusion/lpw.py | 2 +- modules/http/api_router.py | 17 -- modules/http/frontend_router.py | 16 -- modules/http/models/base.py | 9 - modules/http/routes/engine_api.py | 16 -- modules/http/routes/images_api.py | 21 -- modules/http/routes/images_browser.py | 28 --- modules/http/routes/model_api.py | 64 ------ modules/http/routes/plugin_api.py | 32 --- modules/main.py | 15 -- modules/model_manager.py | 94 ++++++-- modules/plugin/plugin_loader.py | 31 --- modules/runners/__init__.py | 105 --------- .../{runner_diffusers.py => diffusers.py} | 97 ++++----- modules/runners/runner.py | 16 +- .../{runner_tensorrt.py => tensorrt.py} | 107 +++++----- modules/shared.py | 11 + modules/tabs/tensorrt.py | 167 +++++++++++++++ modules/tabs/txt2img.py | 85 ++++++++ modules/ui.py | 147 +++++++++++++ modules/utils.py | 9 +- requirements.txt | 26 --- requirements/base.txt | 7 + requirements/dev.txt | 3 + requirements/tensorrt.txt | 7 + script.js | 30 +++ styles.css | 31 +++ webui.py | 32 +++ 49 files changed, 1306 insertions(+), 753 deletions(-) delete mode 100644 api/base.py delete mode 100644 api/models/sockets.py create mode 100644 bin/.gitignore create mode 100644 dev.py create mode 100644 lib/diffusers/scheduler.py delete mode 100644 modules/app.py create mode 100644 modules/components/header.py create mode 100644 modules/components/image_generation_options.py delete mode 100644 modules/http/api_router.py delete mode 100644 modules/http/frontend_router.py delete mode 100644 modules/http/models/base.py delete mode 100644 modules/http/routes/engine_api.py delete mode 100644 modules/http/routes/images_api.py delete mode 100644 modules/http/routes/images_browser.py delete mode 100644 modules/http/routes/model_api.py delete mode 100644 modules/http/routes/plugin_api.py delete mode 100644 modules/main.py delete mode 100644 modules/plugin/plugin_loader.py delete mode 100644 modules/runners/__init__.py rename modules/runners/{runner_diffusers.py => diffusers.py} (58%) rename modules/runners/{runner_tensorrt.py => tensorrt.py} (58%) create mode 100644 modules/shared.py create mode 100644 modules/tabs/tensorrt.py create mode 100644 modules/tabs/txt2img.py create mode 100644 modules/ui.py delete mode 100644 requirements.txt create mode 100644 requirements/base.txt create mode 100644 requirements/dev.txt create mode 100644 requirements/tensorrt.txt create mode 100644 script.js create mode 100644 styles.css create mode 100644 webui.py diff --git a/.gitignore b/.gitignore index 17bf67c1..1a1e84e8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,16 +1,188 @@ -# Python -__pycache__ -/venv +.DS_Store +tmp/ +outputs/ +config.toml dist +node_modules/ + +### Generated by gibo (https://github.com/simonwhitaker/gibo) +### https://raw.github.com/github/gitignore/4488915eec0b3a45b5c63ead28f286819c0917de/Global/VisualStudioCode.gitignore + +.vscode/* +!.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +!.vscode/*.code-snippets + +# Local History for Visual Studio Code +.history/ + +# Built Visual Studio Code Extensions +*.vsix + + +### https://raw.github.com/github/gitignore/4488915eec0b3a45b5c63ead28f286819c0917de/Python.gitignore + +# 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/ + -# Project specific -/config.json -/outputs -/models -/docker-data -/tmp - -# Docs -/docs/.vitepress/dist -/docs/.vitepress/cache -node_modules \ No newline at end of file diff --git a/api/base.py b/api/base.py deleted file mode 100644 index c7de419c..00000000 --- a/api/base.py +++ /dev/null @@ -1,6 +0,0 @@ -from pydantic import BaseModel - - -class BaseModelStream(BaseModel): - def ndjson(self): - return self.json() + "\n" diff --git a/api/models/diffusion.py b/api/models/diffusion.py index ea925102..31ac4610 100644 --- a/api/models/diffusion.py +++ b/api/models/diffusion.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import * from pydantic import BaseModel @@ -18,22 +18,6 @@ class ImageGenerationOptions(BaseModel): img: Optional[str] = None -class ImageGenerationResult(BaseModel): - images: Dict[str, ImageGenerationOptions] - performance: float - - class ImageGenerationError(BaseModel): error: Optional[str] = None message: str - - -class ImageGenerationProgress(BaseModel): - images: Dict[str, ImageGenerationOptions] - progress: float - performance: float - - -class DenoiseLatentData(BaseModel): - step: int - preview: Optional[Dict[str, ImageGenerationOptions]] diff --git a/api/models/plugin.py b/api/models/plugin.py index a3f46dcd..69973602 100644 --- a/api/models/plugin.py +++ b/api/models/plugin.py @@ -13,4 +13,3 @@ class PluginData(BaseModel): meta: PluginMetaData module: str dir: str - js: bool diff --git a/api/models/sockets.py b/api/models/sockets.py deleted file mode 100644 index 1c1d1d12..00000000 --- a/api/models/sockets.py +++ /dev/null @@ -1,10 +0,0 @@ -from typing import Dict, Optional, Union - -from pydantic import BaseModel - - -class SocketData(BaseModel): - namespace: Optional[str] - event: str - id: Optional[str] - data: Union[str, Dict] diff --git a/api/models/tensorrt.py b/api/models/tensorrt.py index 242442e4..2ade45f9 100644 --- a/api/models/tensorrt.py +++ b/api/models/tensorrt.py @@ -28,7 +28,7 @@ class BuildEngineOptions(BaseModel): build_dynamic_shape = True build_all_tactics = False build_preview_features = False - onnx_opset = 16 + onnx_opset = 17 force_engine_build = False force_onnx_export = False force_onnx_optimize = False diff --git a/bin/.gitignore b/bin/.gitignore new file mode 100644 index 00000000..c96a04f0 --- /dev/null +++ b/bin/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore \ No newline at end of file diff --git a/dev.py b/dev.py new file mode 100644 index 00000000..8b6a478d --- /dev/null +++ b/dev.py @@ -0,0 +1,5 @@ +import modules.ui as ui +import webui + +webui.pre_load() +demo = ui.create_ui().queue(64) diff --git a/launch.py b/launch.py index cd11b157..0f05695d 100644 --- a/launch.py +++ b/launch.py @@ -1,24 +1,21 @@ -import argparse import importlib.util import os import platform +import shlex import subprocess import sys -from typing import List, Optional + +commandline_args = os.environ.get("COMMANDLINE_ARGS", "") +sys.argv += shlex.split(commandline_args) python = sys.executable git = os.environ.get("GIT", "git") index_url = os.environ.get("INDEX_URL", "") +stored_commit_hash = None skip_install = False -__dirname__ = os.path.dirname(__file__) -def run( - command: str, - desc: Optional[str] = None, - errdesc: Optional[str] = None, - custom_env: Optional[str] = None, -): +def run(command, desc=None, errdesc=None, custom_env=None): if desc is not None: print(desc) @@ -60,7 +57,14 @@ def is_exe(fpath): return None -def is_installed(package: str): +def check_run(command): + result = subprocess.run( + command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True + ) + return result.returncode == 0 + + +def is_installed(package): try: spec = importlib.util.find_spec(package) except ModuleNotFoundError: @@ -69,7 +73,21 @@ def is_installed(package: str): return spec is not None -def run_pip(args: str, desc: Optional[str] = None): +def commit_hash(): + global stored_commit_hash + + if stored_commit_hash is not None: + return stored_commit_hash + + try: + stored_commit_hash = run(f"{git} rev-parse HEAD").strip() + except Exception: + stored_commit_hash = "" + + return stored_commit_hash + + +def run_pip(args, desc=None): if skip_install: return @@ -81,16 +99,16 @@ def run_pip(args: str, desc: Optional[str] = None): ) -def run_python(code: str, desc: Optional[str] = None, errdesc: Optional[str] = None): +def run_python(code, desc=None, errdesc=None): return run(f'"{python}" -c "{code}"', desc, errdesc) -def extract_arg(args: List[str], name: str): +def extract_arg(args, name): return [x for x in args if x != name], name in args def install_tensorrt(): - trt_version = "8.5.0" + trt_version = "8.6.0" tensorrt_linux_command = os.environ.get( "TENSORRT_LINUX_COMMAND", f"pip install tensorrt=={trt_version}", @@ -121,67 +139,68 @@ def install_tensorrt(): ) -def prepare_environment(args: List[str]): +def prepare_environment(): + commit = commit_hash() + + print(f"Python {sys.version}") + print(f"Commit hash: {commit}") + torch_command = os.environ.get( "TORCH_COMMAND", - "pip install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117", + "pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118", ) - requirements_file = os.environ.get("REQS_FILE", "requirements.txt") - args, skip_install = extract_arg(args, "--skip-install") + sys.argv, skip_install = extract_arg(sys.argv, "--skip-install") if skip_install: return - args, reinstall_torch = extract_arg(args, "--reinstall-torch") - args, reinstall_tensorrt = extract_arg(args, "--reinstall-tensorrt") - - if reinstall_torch or not is_installed("torch"): + sys.argv, reinstall_torch = extract_arg(sys.argv, "--reinstall-torch") + sys.argv, reinstall_xformers = extract_arg(sys.argv, "--reinstall-xformers") + tensorrt = "--tensorrt" in sys.argv + + if reinstall_torch or not is_installed("torch") or not is_installed("torchvision"): + if reinstall_torch: + run( + f'"{python}" -m pip uninstall torch torchvision -y', + "Uninstalling torch and torchvision", + "Couldn't uninstall torch", + ) run( f'"{python}" -m {torch_command}', - "Installing torch", + "Installing torch and torchvision", "Couldn't install torch", ) - run_python( - "import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU'" - ) - - if reinstall_tensorrt or not is_installed("tensorrt"): - install_tensorrt() + if reinstall_xformers or not is_installed("xformers"): + run_pip("install xformers", "xformers") run( - f'"{python}" -m pip install -r {requirements_file}', + f'"{python}" -m pip install -r requirements/base.txt', desc=f"Installing requirements", errdesc=f"Couldn't install requirements", ) + if tensorrt: + run( + f'"{python}" -m pip install -r requirements/tensorrt.txt', + desc=f"Installing tensorrt requirements", + errdesc=f"Couldn't install tensorrt requirements", + ) + install_tensorrt() -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--uvicorn-args", type=str, nargs="*") - ns, args = parser.parse_known_args(sys.argv) - - main_args = os.environ.get("COMMANDLINE_ARGS", "").split(" ") - uvicorn_args = os.environ.get("UVICORN_ARGS", "").split(" ") - - main_args = [*args, *[x for x in main_args if x]] - uvicorn_args = [x for x in uvicorn_args if x] - - if ns.uvicorn_args: - for opt in ns.uvicorn_args: - if "=" in opt: - k, v = opt.split("=") - if type(v) == bool: - if v == True: - uvicorn_args.append(f"--{k}") - else: - uvicorn_args.extend([f"--{k}", v]) - - prepare_environment(main_args) - - env = os.environ.copy() - env["COMMANDLINE_ARGS"] = " ".join(main_args) +def start(): + os.environ["PATH"] = ( + os.path.join(os.path.dirname(__file__), "bin") + + os.pathsep + + os.environ.get("PATH", "") + ) + os.environ["COMMANDLINE_ARGS"] = " ".join(sys.argv[1:]) subprocess.run( - [python, "-m", "uvicorn", "modules.main:sio_app", *uvicorn_args], env=env + [python, "webui.py"], ) + + +if __name__ == "__main__": + prepare_environment() + start() diff --git a/lib/diffusers/scheduler.py b/lib/diffusers/scheduler.py new file mode 100644 index 00000000..8c81e3c4 --- /dev/null +++ b/lib/diffusers/scheduler.py @@ -0,0 +1,14 @@ +import diffusers + +SCHEDULERS = { + "ddim": diffusers.DDIMScheduler, + "deis": diffusers.DEISMultistepScheduler, + "dpm2": diffusers.KDPM2DiscreteScheduler, + "dpm2-a": diffusers.KDPM2AncestralDiscreteScheduler, + "euler_a": diffusers.EulerAncestralDiscreteScheduler, + "euler": diffusers.EulerDiscreteScheduler, + "heun": diffusers.DPMSolverMultistepScheduler, + "dpm++": diffusers.DPMSolverMultistepScheduler, + "dpm": diffusers.DPMSolverMultistepScheduler, + "pndm": diffusers.PNDMScheduler, +} diff --git a/lib/tensorrt/models.py b/lib/tensorrt/models.py index 3ebfacc6..07f7b620 100644 --- a/lib/tensorrt/models.py +++ b/lib/tensorrt/models.py @@ -94,7 +94,7 @@ def __init__( self.embedding_dim = embedding_dim self.text_maxlen = text_maxlen - def get_model(self): + def get_model(self, hf_cache_dir=None): pass def get_input_names(self): @@ -185,9 +185,12 @@ def __init__(self, hf_token, device, path, max_batch_size, embedding_dim): ) self.name = "CLIP" - def get_model(self): + def get_model(self, hf_cache_dir=None): return CLIPTextModel.from_pretrained( - self.path, subfolder="text_encoder", use_auth_token=self.hf_token + self.path, + subfolder="text_encoder", + use_auth_token=self.hf_token, + cache_dir=hf_cache_dir, ).to(self.device) def get_input_names(self): @@ -268,12 +271,13 @@ def __init__( self.unet_dim = unet_dim self.name = "UNet" - def get_model(self): + def get_model(self, hf_cache_dir=None): return UNet2DConditionModel.from_pretrained( self.path, subfolder="unet", use_auth_token=self.hf_token, torch_dtype=torch.float16 if self.fp16 else torch.float32, + cache_dir=hf_cache_dir, ).to(self.device) def get_input_names(self): @@ -372,9 +376,12 @@ def __init__(self, hf_token, device, path, max_batch_size, embedding_dim): ) self.name = "VAE decoder" - def get_model(self): + def get_model(self, hf_cache_dir=None): vae = AutoencoderKL.from_pretrained( - self.path, subfolder="vae", use_auth_token=self.hf_token + self.path, + subfolder="vae", + use_auth_token=self.hf_token, + cache_dir=hf_cache_dir, ).to(self.device) vae.forward = vae.decode return vae @@ -443,11 +450,11 @@ def get_sample_input(self, batch_size, image_height, image_width): class TorchVAEEncoder(torch.nn.Module): - def __init__(self, token, device, path): + def __init__(self, token, device, path, hf_cache_dir=None): super().__init__() self.path = path self.vae_encoder = AutoencoderKL.from_pretrained( - self.path, subfolder="vae", use_auth_token=token + self.path, subfolder="vae", use_auth_token=token, cache_dir=hf_cache_dir ).to(device) def forward(self, x): @@ -465,8 +472,10 @@ def __init__(self, hf_token, device, path, max_batch_size, embedding_dim): ) self.name = "VAE encoder" - def get_model(self): - vae_encoder = TorchVAEEncoder(self.hf_token, self.device, self.path) + def get_model(self, hf_cache_dir=None): + vae_encoder = TorchVAEEncoder( + self.hf_token, self.device, self.path, hf_cache_dir=hf_cache_dir + ) return vae_encoder def get_input_names(self): diff --git a/lib/tensorrt/pipeline.py b/lib/tensorrt/pipeline.py index c39a6f48..e28a710a 100644 --- a/lib/tensorrt/pipeline.py +++ b/lib/tensorrt/pipeline.py @@ -19,7 +19,7 @@ import inspect import os -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import * import diffusers import torch @@ -52,6 +52,7 @@ def from_pretrained( device: Union[str, torch.device], onnx_refit_dir: Optional[str] = None, max_batch_size: int = 1, + hf_cache_dir: Optional[str] = None, ): models = create_models( model_id=model_id, @@ -59,6 +60,7 @@ def from_pretrained( use_auth_token=use_auth_token, device=device, max_batch_size=max_batch_size, + hf_cache_dir=hf_cache_dir, ) engines = {} @@ -91,10 +93,16 @@ def from_pretrained( ) pipe.tokenizer = CLIPTokenizer.from_pretrained( - model_id, subfolder="tokenizer", use_auth_token=use_auth_token + model_id, + subfolder="tokenizer", + use_auth_token=use_auth_token, + cache_dir=hf_cache_dir, ) pipe.scheduler = diffusers.DDIMScheduler.from_pretrained( - model_id, subfolder="scheduler", use_auth_token=use_auth_token + model_id, + subfolder="scheduler", + use_auth_token=use_auth_token, + cache_dir=hf_cache_dir, ) return pipe @@ -161,7 +169,7 @@ def initialize_latents( self, batch_size: int, unet_channels: int, - height: torch.Tensor, + height: int, width: int, dtype: torch.dtype, device: torch.device, @@ -172,6 +180,7 @@ def initialize_latents( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) + shape = (batch_size, unet_channels, height, width) latents = diffusers.utils.randn_tensor( shape, diff --git a/lib/tensorrt/pipeline_stable_diffusion.py b/lib/tensorrt/pipeline_stable_diffusion.py index bf6484b8..90cd6830 100644 --- a/lib/tensorrt/pipeline_stable_diffusion.py +++ b/lib/tensorrt/pipeline_stable_diffusion.py @@ -17,7 +17,7 @@ # limitations under the License. # -from typing import Callable, List, Optional, Union +from typing import * import tensorrt as trt import torch diff --git a/lib/tensorrt/pipeline_stable_diffusion_img2img.py b/lib/tensorrt/pipeline_stable_diffusion_img2img.py index d7bdfec1..41f8cf75 100644 --- a/lib/tensorrt/pipeline_stable_diffusion_img2img.py +++ b/lib/tensorrt/pipeline_stable_diffusion_img2img.py @@ -17,7 +17,7 @@ # limitations under the License. # -from typing import Callable, List, Optional, Union +from typing import * import tensorrt as trt import torch diff --git a/lib/tensorrt/utilities.py b/lib/tensorrt/utilities.py index 2cae3d58..980e7857 100644 --- a/lib/tensorrt/utilities.py +++ b/lib/tensorrt/utilities.py @@ -21,7 +21,7 @@ import gc from collections import OrderedDict from copy import copy -from typing import Dict, List, Optional, Union +from typing import * import diffusers import numpy as np @@ -343,12 +343,19 @@ def create_models( use_auth_token: Optional[str], device: Union[str, torch.device], max_batch_size: int, + hf_cache_dir: Optional[str] = None, ): text_encoder_config = CLIPTextConfig.from_pretrained( - model_id, subfolder="text_encoder", use_auth_token=use_auth_token + model_id, + subfolder="text_encoder", + use_auth_token=use_auth_token, + cache_dir=hf_cache_dir, ) unet = diffusers.UNet2DConditionModel.from_pretrained( - model_id, subfolder="unet", use_auth_token=use_auth_token + model_id, + subfolder="unet", + use_auth_token=use_auth_token, + cache_dir=hf_cache_dir, ) unet_dim = unet.in_channels embedding_dim = text_encoder_config.hidden_size @@ -438,8 +445,9 @@ def export_onnx( opt_image_height: int, opt_image_width: int, onnx_opset: int, + hf_cache_dir: Optional[str] = None, ): - model = model_data.get_model() + model = model_data.get_model(hf_cache_dir=hf_cache_dir) with torch.inference_mode(), torch.autocast("cuda"): inputs = model_data.get_sample_input(1, opt_image_height, opt_image_width) torch.onnx.export( @@ -465,3 +473,6 @@ def optimize_onnx( ): onnx_opt_graph = model_data.optimize(onnx.load(onnx_path)) onnx.save(onnx_opt_graph, onnx_opt_path) + del onnx_opt_graph + torch.cuda.empty_cache() + gc.collect() diff --git a/modules/acceleration/tensorrt/engine.py b/modules/acceleration/tensorrt/engine.py index 5ac0fdcb..671662fa 100644 --- a/modules/acceleration/tensorrt/engine.py +++ b/modules/acceleration/tensorrt/engine.py @@ -1,14 +1,10 @@ import gc import os -import onnx import torch -from cuda import cudart from api.models.tensorrt import BuildEngineOptions, TensorRTEngineData -from lib.tensorrt.models import BaseModel from lib.tensorrt.utilities import ( - Engine, build_engine, create_models, export_onnx, @@ -16,6 +12,7 @@ ) from modules import model_manager from modules.logger import logger +from modules.shared import hf_cache_dir def create_onnx_path(name, onnx_dir, opt=True): @@ -33,81 +30,9 @@ def __init__(self, opts: BuildEngineOptions): device=torch.device("cuda"), use_auth_token=opts.hf_token, max_batch_size=opts.max_batch_size, + hf_cache_dir=hf_cache_dir(), ) - def build_onnx( - self, - model_name: str, - onnx_dir: str, - model_data: BaseModel, - ): - onnx_path = create_onnx_path(model_name, onnx_dir, opt=False) - onnx_opt_path = create_onnx_path(model_name, onnx_dir) - if self.opts.force_onnx_export or not os.path.exists(onnx_path): - logger.info(f"Exporting model: {onnx_path}") - model = model_data.get_model() - with torch.inference_mode(), torch.autocast("cuda"): - inputs = model_data.get_sample_input( - 1, self.opts.opt_image_height, self.opts.opt_image_width - ) - torch.onnx.export( - model, - inputs, - onnx_path, - export_params=True, - opset_version=self.opts.onnx_opset, - do_constant_folding=True, - input_names=model_data.get_input_names(), - output_names=model_data.get_output_names(), - dynamic_axes=model_data.get_dynamic_axes(), - ) - del model - torch.cuda.empty_cache() - gc.collect() - else: - logger.info(f"Found cached model: {onnx_path}") - - if self.opts.force_onnx_optimize or not os.path.exists(onnx_opt_path): - logger.info(f"Generating optimizing model: {onnx_opt_path}") - onnx_opt_graph = model_data.optimize(onnx.load(onnx_path)) - onnx.save(onnx_opt_graph, onnx_opt_path) - - def build_engine( - self, - model_name: str, - engine_dir: str, - onnx_dir: str, - model_data: BaseModel, - ): - _, free_mem, _ = cudart.cudaMemGetInfo() - GiB = 2**30 - if free_mem > 6 * GiB: - activation_carveout = 4 * GiB - max_workspace_size = free_mem - activation_carveout - else: - max_workspace_size = 0 - model_data.min_latent_shape = self.opts.min_latent_resolution // 8 - model_data.max_latent_shape = self.opts.max_latent_resolution // 8 - engine = Engine(os.path.join(engine_dir, f"{model_name}.plan")) - onnx_opt_path = create_onnx_path(model_name, onnx_dir) - engine.build( - onnx_opt_path, - fp16=True, - input_profile=model_data.get_input_profile( - 1, - self.opts.opt_image_height, - self.opts.opt_image_width, - static_batch=self.opts.build_static_batch, - static_shape=not self.opts.build_dynamic_shape, - ), - enable_all_tactics=self.opts.build_all_tactics, - enable_refit=self.opts.build_enable_refit, - enable_preview=self.opts.build_preview_features, - workspace_size=max_workspace_size, - ) - - return engine - def build(self): model_dir = self.model.get_trt_path() engine_dir = os.path.join(model_dir, "engine") @@ -127,6 +52,7 @@ def build(self): opt_image_height=self.opts.opt_image_height, opt_image_width=self.opts.opt_image_width, onnx_opset=self.opts.onnx_opset, + hf_cache_dir=hf_cache_dir(), ) if not self.opts.force_onnx_optimize and os.path.exists(onnx_opt_path): logger.info(f"Found cached model: {onnx_opt_path}") diff --git a/modules/acceleration/tensorrt/text_encoder.py b/modules/acceleration/tensorrt/text_encoder.py index 7c504900..fce77ea6 100644 --- a/modules/acceleration/tensorrt/text_encoder.py +++ b/modules/acceleration/tensorrt/text_encoder.py @@ -1,4 +1,3 @@ -import numpy as np import torch from polygraphy import cuda diff --git a/modules/app.py b/modules/app.py deleted file mode 100644 index 98808752..00000000 --- a/modules/app.py +++ /dev/null @@ -1,35 +0,0 @@ -import mimetypes - -import socketio -from fastapi import FastAPI -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import RedirectResponse -from fastapi.routing import APIRoute - -from . import config - - -def custom_generate_unique_id(route: APIRoute): - return route.name - - -app = FastAPI(generate_unique_id_function=custom_generate_unique_id) -allowed_hosts = config.get("allow_hosts") -app.add_middleware( - CORSMiddleware, - allow_origins=allowed_hosts.split(",") if allowed_hosts else [], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) -mimetypes.init() -mimetypes.add_type("application/javascript", ".js") - - -@app.get("/") -def redirect(): - return RedirectResponse("/app") - - -sio = socketio.AsyncServer(async_mode="asgi") -sio_app = socketio.ASGIApp(sio, other_asgi_app=app) diff --git a/modules/components/header.py b/modules/components/header.py new file mode 100644 index 00000000..057718a3 --- /dev/null +++ b/modules/components/header.py @@ -0,0 +1,140 @@ +import gradio as gr + +from modules import model_manager + + +def model_list_str(): + return [x.model_id for x in model_manager.sd_models] + + +def change_model(model_id: str): + if model_id not in model_list_str(): + raise ValueError("Model not found.") + model_manager.set_model(model_id) + return model_id + + +def add_model(model_id: str): + if model_id not in model_list_str(): + searched = model_manager.search_model(model_id) + if len(searched) < 1: + raise ValueError("Model not found.") + model_manager.add_model(model_id) + return gr.Dropdown.update(choices=model_list_str()) + + +def diffusers_ui(): + model_id = ( + model_manager.sd_model.model_id if model_manager.sd_model is not None else None + ) + with gr.Row(visible=model_manager.mode == "diffusers") as row: + with gr.Column(scale=0.25): + with gr.Row(): + model_id_dropdown = gr.Dropdown( + value=model_id, + choices=model_list_str(), + show_label=False, + ) + reload_models_button = gr.Button("🔄", elem_classes=["tool-button"]) + + with gr.Column(scale=0.25): + with gr.Row(): + add_model_textbox = gr.Textbox( + placeholder="Add model", + show_label=False, + ) + add_model_button = gr.Button("💾", elem_classes=["tool-button"]) + + model_id_dropdown.change( + fn=change_model, inputs=[model_id_dropdown], outputs=[model_id_dropdown] + ) + reload_models_button.click( + fn=lambda: gr.Dropdown.update( + choices=model_list_str(), + value=model_manager.sd_model.model_id + if model_manager.sd_model is not None + else None, + ), + inputs=[], + outputs=[model_id_dropdown], + ) + add_model_button.click( + fn=add_model, inputs=[add_model_textbox], outputs=[model_id_dropdown] + ) + + return row + + +def tensorrt_ui(): + def tensorrt_models(): + return [x for x in model_manager.sd_models if x.trt_available()] + + model_id = ( + model_manager.sd_model.model_id + if model_manager.sd_model is not None and model_manager.sd_model.trt_available() + else None + ) + with gr.Row(visible=model_manager.mode == "tensorrt") as row: + with gr.Column(scale=0.25): + with gr.Row(): + model_id_dropdown = gr.Dropdown( + value=model_id, + choices=[x.model_id for x in tensorrt_models()], + show_label=False, + ) + reload_models_button = gr.Button("🔄", elem_classes=["tool-button"]) + + reload_models_button.click( + fn=lambda: gr.Dropdown.update( + choices=[x.model_id for x in tensorrt_models()], + value=model_manager.sd_model.model_id + if model_manager.sd_model is not None + else None, + ), + inputs=[], + outputs=[model_id_dropdown], + ) + model_id_dropdown.change( + fn=change_model, + inputs=[model_id_dropdown], + outputs=[model_id_dropdown], + ) + + return row + + +def ui(): + with gr.Box(): + with gr.Row(): + diffusers = diffusers_ui() + tensorrt = tensorrt_ui() + with gr.Row(): + mode = gr.Radio( + choices=model_manager.available_mode, + value=model_manager.mode, + show_label=False, + ) + reload_button = gr.Button( + "🔄", + elem_classes=["tool-button"], + elem_id="inference-mode-reload-button", + ) + + def reload(): + return ( + model_manager.mode, + gr.Row.update(visible=model_manager.mode == "diffusers"), + gr.Row.update(visible=model_manager.mode == "tensorrt"), + ) + + def on_change(mode: str): + model_manager.set_mode(mode) + return ( + mode, + gr.Row.update(visible=mode == "diffusers"), + gr.Row.update(visible=mode == "tensorrt"), + ) + + reload_button.click(fn=reload, inputs=[], outputs=[mode, diffusers, tensorrt]) + + mode.change(fn=on_change, inputs=[mode], outputs=[mode, diffusers, tensorrt]) diff --git a/modules/components/image_generation_options.py b/modules/components/image_generation_options.py new file mode 100644 index 00000000..d2deae4d --- /dev/null +++ b/modules/components/image_generation_options.py @@ -0,0 +1,94 @@ +import gradio as gr + +from lib.diffusers.scheduler import SCHEDULERS + + +def ui(): + with gr.Column(): + with gr.Row(): + with gr.Column(scale=3): + prompt_textbox = gr.TextArea( + "", + lines=3, + placeholder="Prompt", + show_label=False, + ) + negative_prompt_textbox = gr.TextArea( + "", + lines=3, + placeholder="Negative Prompt", + show_label=False, + ) + generate_button = gr.Button( + "Generate", + variant="primary", + ) + with gr.Row(): + with gr.Column(scale=1.25): + with gr.Row(): + sampler_dropdown = gr.Dropdown( + choices=list(SCHEDULERS.keys()), + value="euler_a", + label="Sampler", + ) + sampling_steps_slider = gr.Slider( + value=25, + minimum=1, + maximum=100, + step=1, + label="Sampling Steps", + ) + with gr.Row(): + batch_size_slider = gr.Slider( + value=1, + minimum=1, + maximum=50, + step=1, + label="Batch size", + ) + batch_count_slider = gr.Slider( + value=1, + minimum=1, + maximum=50, + step=1, + label="Batch count", + ) + with gr.Row(): + cfg_scale_slider = gr.Slider( + value=7.5, + minimum=1, + maximum=20, + step=0.5, + label="CFG Scale", + ) + with gr.Row(): + width_slider = gr.Slider( + value=512, minimum=128, maximum=2048, step=64, label="Width" + ) + height_slider = gr.Slider( + value=512, minimum=128, maximum=2048, step=64, label="Height" + ) + with gr.Row(): + seed_number = gr.Number( + value=-1, + ) + with gr.Column(): + output_images = gr.Gallery( + elem_classes="image_generation_gallery" + ).style(columns=4) + status_textbox = gr.Textbox(interactive=False, show_label=False) + + prompts = [prompt_textbox, negative_prompt_textbox] + options = [ + sampler_dropdown, # sampler name + sampling_steps_slider, # num sampling steps + batch_size_slider, # batch size + batch_count_slider, # batch count + cfg_scale_slider, # cfg scale + width_slider, # width + height_slider, # height + seed_number, # seed + ] + outputs = [output_images, status_textbox] + + return generate_button, prompts, options, outputs diff --git a/modules/config.py b/modules/config.py index ddf4399b..ccddfdf8 100644 --- a/modules/config.py +++ b/modules/config.py @@ -1,7 +1,8 @@ import argparse -import json import os +import toml + DEFAULT_CONFIG = { "version": "0", "images/txt2img/save_dir": "outputs/txt2img", @@ -19,12 +20,18 @@ parser = argparse.ArgumentParser() -parser.add_argument("--allow-hosts", type=str, default="") -parser.add_argument("--xformers", action="store_true") +parser.add_argument("--config-file", type=str, default="config.toml") + +parser.add_argument("--host", type=str, default="") +parser.add_argument("--port", type=int, default=7860) +parser.add_argument("--share", action="store_true") + parser.add_argument("--model-dir", type=str, default="models") -parser.add_argument("--config-file", type=str, default="config.json") parser.add_argument("--hf-token", type=str) +parser.add_argument("--xformers", action="store_true") +parser.add_argument("--tensorrt", action="store_true") + cmd_opts, _ = parser.parse_known_args( os.environ["COMMANDLINE_ARGS"].split(" ") if "COMMANDLINE_ARGS" in os.environ @@ -36,19 +43,24 @@ def get_config(): + if not os.path.exists(cmd_opts.config_file): + with open(cmd_opts.config_file, "w") as f: + f.write(toml.dumps(DEFAULT_CONFIG)) + with open(cmd_opts.config_file, mode="r") as f: txt = f.read() try: - config = json.loads(txt) - except: + config = toml.loads(txt) + except Exception as e: + print(e) config = DEFAULT_CONFIG return config def save_config(): with open(cmd_opts.config_file, mode="w") as f: - f.write(json.dumps(opts)) + f.write(toml.dumps(opts)) def set(key: str, value: str): diff --git a/modules/diffusion/lpw.py b/modules/diffusion/lpw.py index 99fe2479..b7924caa 100644 --- a/modules/diffusion/lpw.py +++ b/modules/diffusion/lpw.py @@ -1,5 +1,5 @@ import re -from typing import List, Optional, Union +from typing import * import torch from transformers import CLIPTokenizer diff --git a/modules/http/api_router.py b/modules/http/api_router.py deleted file mode 100644 index 5fed566c..00000000 --- a/modules/http/api_router.py +++ /dev/null @@ -1,17 +0,0 @@ -import importlib -import os - -from fastapi import APIRouter - -__dirname__ = os.path.dirname(__file__) - -api = APIRouter(prefix="/api", tags=["main"]) - - -def initialize(): - for route in os.listdir(os.path.join(__dirname__, "routes")): - module_name = "modules.http.routes." + route.replace(".py", "") - importlib.import_module(module_name) - - -initialize() diff --git a/modules/http/frontend_router.py b/modules/http/frontend_router.py deleted file mode 100644 index 3d6d510e..00000000 --- a/modules/http/frontend_router.py +++ /dev/null @@ -1,16 +0,0 @@ -import os - -from fastapi import APIRouter, Request -from fastapi.responses import FileResponse - -from ..config import ROOT_DIR - -frontend = APIRouter(prefix="/app", tags=["application"]) - - -@frontend.get("/{full_path:path}") -def handler(_: Request, full_path: str): - full_path = full_path.replace("/", os.sep) - if full_path == "": - full_path = "index.html" - return FileResponse(os.path.join(ROOT_DIR, "frontend", "app", full_path)) diff --git a/modules/http/models/base.py b/modules/http/models/base.py deleted file mode 100644 index 33a13ae8..00000000 --- a/modules/http/models/base.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import Any, Optional - -from pydantic import BaseModel - - -class BaseResponseModel(BaseModel): - status: str - message: Optional[str] - data: Any diff --git a/modules/http/routes/engine_api.py b/modules/http/routes/engine_api.py deleted file mode 100644 index 5fac4427..00000000 --- a/modules/http/routes/engine_api.py +++ /dev/null @@ -1,16 +0,0 @@ -from api.models.tensorrt import BuildEngineOptions -from modules import model_manager -from modules.http.models.base import BaseResponseModel - -from ...acceleration.tensorrt.engine import EngineBuilder -from ..api_router import api - - -@api.post("/engine/build") -async def build_engine(req: BuildEngineOptions): - model_manager.runner.teardown() - builder = EngineBuilder(req) - builder.build() - model_manager.runner.activate() - - return BaseResponseModel(status="success", message="Finish build engine") diff --git a/modules/http/routes/images_api.py b/modules/http/routes/images_api.py deleted file mode 100644 index bcbf434d..00000000 --- a/modules/http/routes/images_api.py +++ /dev/null @@ -1,21 +0,0 @@ -import asyncio - -from api.models.diffusion import ImageGenerationOptions, ImageGenerationResult -from modules import model_manager - -from ..api_router import api -from ..models.base import BaseResponseModel - - -class GenerateImageResponseModel(BaseResponseModel): - data: ImageGenerationResult - - -@api.post("/images/generate", response_model=GenerateImageResponseModel) -async def generate_image(opts: ImageGenerationOptions): - loop = asyncio.get_event_loop() - result = await loop.run_in_executor(None, model_manager.runner.generate, opts) - return GenerateImageResponseModel( - status="success", - data=result, - ) diff --git a/modules/http/routes/images_browser.py b/modules/http/routes/images_browser.py deleted file mode 100644 index c8e5d667..00000000 --- a/modules/http/routes/images_browser.py +++ /dev/null @@ -1,28 +0,0 @@ -from fastapi.responses import FileResponse - -from modules import images - -from ..api_router import api -from ..models.base import BaseResponseModel - - -class FileListResponseModel(BaseResponseModel): - length: int - data: dict[str, dict] = {} - - -@api.get("/images/browser/{category}/{page}", response_model=FileListResponseModel) -def get_all_image_files(category: str, page: int): - count = 20 - files = images.get_all_image_files(category) - data = files[page * count : page * count + count] - data = { - d: {"info": images.get_image_parameter(images.get_image(category, d))} - for d in data - } - return FileListResponseModel(status="success", length=len(files), data=data) - - -@api.get("/images/browser/get/{category}/{filename}") -def get_image(category: str, filename: str): - return FileResponse(images.get_image_filepath(category, filename)) diff --git a/modules/http/routes/model_api.py b/modules/http/routes/model_api.py deleted file mode 100644 index 00ce5c9a..00000000 --- a/modules/http/routes/model_api.py +++ /dev/null @@ -1,64 +0,0 @@ -from typing import List, Optional - -from pydantic import BaseModel - -from modules import config, model_manager -from modules.model import StableDiffusionModel -from modules.model_manager import ModelMode - -from ..api_router import api -from ..models.base import BaseResponseModel - - -class ModelListResponseModel(BaseResponseModel): - data: List[StableDiffusionModel] = [] - - -@api.get("/model/list", response_model=ModelListResponseModel) -def get_models(): - return ModelListResponseModel(status="success", data=model_manager.sd_models) - - -class ModelCurrentResponseModel(BaseResponseModel): - data: Optional[StableDiffusionModel] - - -@api.get("/model/currnet", response_model=ModelCurrentResponseModel) -def get_current_model(): - return ModelCurrentResponseModel(status="success", data=model_manager.sd_model) - - -class SetModelRequest(BaseModel): - model_id: str - - -@api.post("/model/current", response_model=BaseResponseModel) -def set_model(req: SetModelRequest): - model_manager.set_model(req.model_id) - config.set("model", req.model_id) - return BaseResponseModel(status="success", message="Set model") - - -class ModelModeResponseModel(BaseResponseModel): - data: ModelMode - - -@api.get("/model/mode", response_model=ModelModeResponseModel) -def get_model_mode(): - return ModelModeResponseModel(status="success", data=model_manager.mode) - - -class SetModeRequest(BaseModel): - mode: ModelMode - - -@api.post("/model/mode", response_model=BaseResponseModel) -def set_model_mode(req: SetModeRequest): - model_manager.set_mode(req.mode) - config.set("mode", req.mode) - return BaseResponseModel(status="success", message="Set mode") - - -@api.get("/model/{model_id}/trt_available") -def trt_available(model_id: str): - return model_manager.get_model(model_id).trt_available() diff --git a/modules/http/routes/plugin_api.py b/modules/http/routes/plugin_api.py deleted file mode 100644 index 49bbf63d..00000000 --- a/modules/http/routes/plugin_api.py +++ /dev/null @@ -1,32 +0,0 @@ -import os -from typing import List - -from fastapi.exceptions import HTTPException -from fastapi.responses import FileResponse - -from modules.plugin.plugin_loader import PluginData, plugins - -from ..api_router import api -from ..models.base import BaseResponseModel - - -class PluginListResponseModel(BaseResponseModel): - data: List[PluginData] - - -@api.get("/plugins/list", response_model=PluginListResponseModel) -def plugin_list(): - return PluginListResponseModel(status="success", data=plugins) - - -@api.get("/plugins/js/{plugin_name}", response_model=PluginListResponseModel) -def plugin(plugin_name: str): - plugin = [x for x in plugins if x.meta.name == plugin_name] - if len(plugin) < 1: - raise HTTPException(status_code=404, detail="Plugin not found") - else: - plugin = plugin[0] - - return FileResponse( - os.path.join(plugin.dir, "main.js"), media_type="application/javascript" - ) diff --git a/modules/main.py b/modules/main.py deleted file mode 100644 index dfb765d0..00000000 --- a/modules/main.py +++ /dev/null @@ -1,15 +0,0 @@ -from modules.app import app, sio_app - -from . import config, model_manager -from .http.api_router import api -from .http.frontend_router import frontend -from .plugin import plugin_loader - -sio_app - -app.include_router(api) -app.include_router(frontend) - -config.init() -model_manager.set_default_model() -plugin_loader.load_plugins() diff --git a/modules/model_manager.py b/modules/model_manager.py index e084aa63..7f82d812 100644 --- a/modules/model_manager.py +++ b/modules/model_manager.py @@ -1,43 +1,52 @@ -from typing import List, Literal +from typing import * + +from huggingface_hub import HfApi, ModelFilter from modules.logger import logger -from modules.runners.runner_diffusers import DiffusersDiffusionRunner from . import config from .model import StableDiffusionModel -from .runners.runner_tensorrt import TensorRTDiffusionRunner +from .utils import tensorrt_is_available ModelMode = Literal["diffusers", "tensorrt"] runner = None mode: ModelMode = config.get("mode") sd_models: List[StableDiffusionModel] = [] -sd_model: StableDiffusionModel = None +sd_model: Optional[StableDiffusionModel] = None +available_mode = ["diffusers"] -raw_model_list = config.get("models") or [] -if len(raw_model_list) < 1: - raw_model_list = config.DEFAULT_CONFIG["models"] -for model_data in raw_model_list: - sd_models.append(StableDiffusionModel(**model_data)) +def init(): + global mode + raw_model_list = config.get("models") or [] + if len(raw_model_list) < 1: + raw_model_list = config.DEFAULT_CONFIG["models"] + for model_data in raw_model_list: + sd_models.append(StableDiffusionModel(**model_data)) -def set_default_model(): - global sd_model - prev = config.get("model") - sd_model = [x for x in sd_models if x.model_id == prev] - if len(sd_model) != 1: - sd_model = sd_models[0] - else: - sd_model = sd_model[0] + trt_module_status, trt_version_status = tensorrt_is_available() + if config.get("tensorrt"): + if trt_module_status and trt_version_status: + available_mode.append("tensorrt") + elif trt_module_status: + logger.warning( + "TensorRT is available, but torch version is not compatible." + ) - set_model(sd_model.model_id) + if mode not in available_mode: + mode = available_mode[0] + config.set("mode", mode) def set_mode(m: ModelMode): global mode + if m == mode: + return mode = m runner.teardown() set_model(sd_model.model_id) + config.set("mode", mode) def get_model(model_id: str): @@ -47,6 +56,12 @@ def get_model(model_id: str): return model[0] +def add_model(model_id: str): + global sd_models + sd_models.append(StableDiffusionModel(model_id=model_id)) + config.set("models", [x.dict() for x in sd_models]) + + def set_model(model_id: str): global runner global sd_model @@ -56,9 +71,52 @@ def set_model(model_id: str): else: sd_model = sd_model[0] + if runner is not None: + runner.teardown() + del runner + logger.info(f"Loading {sd_model.model_id}...") if mode == "diffusers": + from modules.runners.diffusers import DiffusersDiffusionRunner + runner = DiffusersDiffusionRunner(sd_model) elif mode == "tensorrt": + from .runners.tensorrt import TensorRTDiffusionRunner + runner = TensorRTDiffusionRunner(sd_model) logger.info(f"Loaded {sd_model.model_id}...") + + config.set("model", sd_model.model_id) + + +def set_default_model(): + global sd_model + prev = config.get("model") + sd_model = [x for x in sd_models if x.model_id == prev] + + if len(sd_model) == 1: + sd_model = sd_model[0] + + if mode == "tensorrt" and not sd_model.trt_available(): + sd_model = None + + if sd_model is None: + available_models = [*sd_models] + + if mode == "tensorrt": + available_models = [x for x in available_models if x.trt_available()] + + if len(available_models) < 1: + set_mode("diffusers") + + sd_model = available_models[0] + + set_model(sd_model.model_id) + + +def search_model(model_id: str): + api = HfApi() + models = api.list_models( + filter=ModelFilter(library="diffusers", model_name=model_id) + ) + return models diff --git a/modules/plugin/plugin_loader.py b/modules/plugin/plugin_loader.py deleted file mode 100644 index 495a2076..00000000 --- a/modules/plugin/plugin_loader.py +++ /dev/null @@ -1,31 +0,0 @@ -import importlib -import os -from typing import List - -from api.models.plugin import PluginData, PluginMetaData -from modules import config - -plugins: List[PluginData] = [] - - -def load_plugins(): - plugin_dir = os.path.join(config.ROOT_DIR, "plugins") - os.makedirs(plugin_dir, exist_ok=True) - for dir in os.listdir(plugin_dir): - fullpath = os.path.join(plugin_dir, dir) - meta_path = os.path.join(fullpath, "plugin.json") - if not os.path.exists(meta_path): - continue - meta = PluginMetaData.parse_file(meta_path) - main_module = f"plugins.{dir}.{meta.main}" - try: - importlib.import_module(main_module, meta.name) - data = PluginData( - meta=meta, - module=main_module, - dir=fullpath, - js=os.path.exists(os.path.join(fullpath, "main.js")), - ) - plugins.append(data) - except Exception as e: - print(f"Failed to load plugin: {meta.name}", e) diff --git a/modules/runners/__init__.py b/modules/runners/__init__.py deleted file mode 100644 index f87c5340..00000000 --- a/modules/runners/__init__.py +++ /dev/null @@ -1,105 +0,0 @@ -# import glob -# import json -# import os -# from typing import Literal - -# from api.models.diffusion import ImageGenerationOptions -# from modules import config, utils -# from modules.images import save_image -# from modules.diffusion.runner import BaseRunner -# from modules.diffusion.tensorrt.runner import TensorRTDiffusionRunner - -# current: BaseRunner = None -# mode: Literal["diffusers", "tensorrt"] = "diffusers" - - -# def meta_filename(): -# if mode == "diffusers": -# return "model_index.json" -# elif mode == "tensorrt": -# return "tensorrt.json" - - -# def set_runner( -# model_dir: str, -# tokenizer_id="openai/clip-vit-large-patch14", -# ): -# global current -# if current is not None: -# current.teardown() - -# meta_path = os.path.join(config.get("model_dir"), model_dir, meta_filename()) - -# try: -# with open(meta_path, mode="r") as f: -# meta = json.loads(f.read()) -# with open(meta_path, mode="w") as f: -# f.write(json.dumps(meta)) -# except Exception as e: -# print(e) - -# try: -# current = TensorRTDiffusionRunner( -# os.path.join(config.get("model_dir"), model_dir).replace("/", os.sep) -# ) -# current.activate(tokenizer_id) -# config.set("model", model_dir) -# except RuntimeError: -# print(f"Failed to load model: {model_dir}") - - -# def get_runners(): -# model_dirs = glob.glob( -# os.path.join(config.get("model_dir"), "**", meta_filename()), -# recursive=True, -# ) - -# return [ -# os.path.relpath(os.path.dirname(x), config.get("model_dir")).replace( -# os.sep, "/" -# ) -# for x in model_dirs -# ] - - -# def generate(options: ImageGenerationOptions): -# gen = current.infer(options) -# result = next(gen) -# gen.close() - -# for img, info in result.images.items(): -# save_image(utils.b642img(img), info) - -# return result - - -# def generator(options: ImageGenerationOptions): -# options.generator = True -# for data in current.infer(options): -# yield data -# if data.type == "result": -# for img, info in data.images.items(): -# save_image(utils.b642img(img), info) - - -# def set_default_model(): -# if os.path.exists(config.get("model_dir")): -# models = glob.glob( -# os.path.join(config.get("model_dir"), "**", meta_filename()), -# recursive=True, -# ) - -# if len(models) < 1: -# return - -# previous = config.get("model") -# if previous is not None: -# model = os.path.join(config.get("model_dir"), previous) -# if os.path.exists(model): -# set_runner(previous) -# return - -# model_dir = os.path.relpath(os.path.dirname(models[0]), config.get("model_dir")) - -# set_runner(model_dir) -# config.set("model", model_dir) diff --git a/modules/runners/runner_diffusers.py b/modules/runners/diffusers.py similarity index 58% rename from modules/runners/runner_diffusers.py rename to modules/runners/diffusers.py index a9dc73c5..14046105 100644 --- a/modules/runners/runner_diffusers.py +++ b/modules/runners/diffusers.py @@ -1,22 +1,18 @@ import gc import random -import time -from itertools import chain -from typing import List, Optional, Union +from concurrent.futures import ThreadPoolExecutor +from queue import Queue +from typing import * import torch from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionPipeline -from api.models.diffusion import ( - DenoiseLatentData, - ImageGenerationOptions, - ImageGenerationResult, -) +from api.models.diffusion import ImageGenerationOptions from modules import config, utils -from modules.app import sio from modules.diffusion.lpw import LongPromptWeightingPipeline from modules.images import save_image from modules.model import StableDiffusionModel +from modules.shared import hf_cache_dir from .runner import BaseRunner @@ -34,10 +30,11 @@ def activate(self) -> None: self.model.model_id, use_auth_token=config.get("hf_token"), torch_dtype=torch.float16, - custom_pipeline="lpw_stable_diffusion", + cache_dir=hf_cache_dir(), ).to( torch.device("cuda") ) + self.pipe.safety_checker = None self.pipe.enable_attention_slicing() if utils.is_installed("xformers") and config.get("xformers"): @@ -61,59 +58,65 @@ def _encode_prompt( self.pipe._encode_prompt = _encode_prompt - def teardown(self) -> None: - del self.pipe + def teardown(self): + if hasattr(self, "pipe"): + del self.pipe torch.cuda.empty_cache() gc.collect() - def generate(self, opts: ImageGenerationOptions) -> ImageGenerationResult: + def generate(self, opts: ImageGenerationOptions): self.wait_loading() + + results = [] + if opts.seed is None or opts.seed == -1: opts.seed = random.randrange(0, 4294967294, 1) - e2e_tic = time.perf_counter() - results = [] + self.pipe.scheduler = self.get_scheduler(opts.scheduler_id) for i in range(opts.batch_count): manual_seed = opts.seed + i + generator = torch.Generator(device=self.pipe.device).manual_seed( + manual_seed + ) + + queue = Queue() + done = object() + def callback( step: int, timestep: torch.Tensor, latents: torch.Tensor, ): - timesteps = self.pipe.scheduler.timesteps - include = step % 10 == 0 and len(timesteps) - step >= 10 - if include: - images = self.pipe.decode_latents(latents) - images = self.pipe.numpy_to_pil(images) - images = [ - *images, - *list(chain.from_iterable([x for x, _ in results])), - ] - data = DenoiseLatentData( - step=step, - preview={utils.img2b64(x): opts for x in images} if include else {}, + queue.put(((opts.steps * i) + step, results)) + + def on_done(feature): + queue.put(done) + + with ThreadPoolExecutor() as executer: + feature = executer.submit( + self.pipe, + prompt=[opts.prompt] * opts.batch_size, + negative_prompt=[opts.negative_prompt] * opts.batch_size, + height=opts.image_height, + width=opts.image_width, + guidance_scale=opts.scale, + num_inference_steps=opts.steps, + generator=generator, + callback=callback, ) + feature.add_done_callback(on_done) - async def runner(): - await sio.emit("denoise_latent", data=data.dict()) + while True: + data = queue.get() + if data is done: + break + else: + yield data - utils.fire_and_forget(runner)() + images = feature.result().images - generator = torch.Generator(device=self.pipe.device).manual_seed( - manual_seed - ) - images = self.pipe( - prompt=opts.prompt, - negative_prompt=opts.negative_prompt, - height=opts.image_height, - width=opts.image_width, - guidance_scale=opts.scale, - num_inference_steps=opts.steps, - generator=generator, - callback=callback, - ).images results.append( ( images, @@ -125,10 +128,4 @@ async def runner(): for img in images: save_image(img, opts) - all_perf_time = time.perf_counter() - e2e_tic - result = ImageGenerationResult(images={}, performance=all_perf_time) - for images, opts in results: - for img in images: - result.images[utils.img2b64(img)] = opts - - return result + yield results diff --git a/modules/runners/runner.py b/modules/runners/runner.py index e798ad6e..938e6731 100644 --- a/modules/runners/runner.py +++ b/modules/runners/runner.py @@ -1,7 +1,11 @@ import time +from typing import * -from api.models.diffusion import ImageGenerationOptions, ImageGenerationResult +from api.models.diffusion import ImageGenerationOptions +from lib.diffusers.scheduler import SCHEDULERS +from modules import config from modules.model import StableDiffusionModel +from modules.shared import hf_cache_dir class BaseRunner: @@ -19,5 +23,13 @@ def wait_loading(self): while self.loading: time.sleep(0.5) - def generate(self, opts: ImageGenerationOptions) -> ImageGenerationResult: + def generate(self, opts: ImageGenerationOptions): pass + + def get_scheduler(self, scheduler_name: str): + return SCHEDULERS[scheduler_name].from_pretrained( + self.model.model_id, + subfolder="scheduler", + use_auth_token=config.get("hf_token"), + cache_dir=hf_cache_dir(), + ) diff --git a/modules/runners/runner_tensorrt.py b/modules/runners/tensorrt.py similarity index 58% rename from modules/runners/runner_tensorrt.py rename to modules/runners/tensorrt.py index ccf75ed1..36359194 100644 --- a/modules/runners/runner_tensorrt.py +++ b/modules/runners/tensorrt.py @@ -1,29 +1,24 @@ import gc import os import random -import time -from itertools import chain +from concurrent.futures import ThreadPoolExecutor +from queue import Queue +from typing import * from typing import Optional, Union -import numpy as np import torch -from polygraphy import cuda -from api.models.diffusion import ( - DenoiseLatentData, - ImageGenerationOptions, - ImageGenerationResult, -) +from api.models.diffusion import ImageGenerationOptions from lib.tensorrt.pipeline_stable_diffusion import TensorRTStableDiffusionPipeline from lib.tensorrt.pipeline_stable_diffusion_img2img import ( TensorRTStableDiffusionImg2ImgPipeline, ) -from modules import config, utils +from modules import config from modules.acceleration.tensorrt.text_encoder import TensorRTCLIPTextModel -from modules.app import sio from modules.diffusion.lpw import LongPromptWeightingPipeline from modules.images import save_image from modules.model import StableDiffusionModel +from modules.shared import hf_cache_dir from .runner import BaseRunner @@ -49,6 +44,7 @@ def activate(self): use_auth_token=config.get("hf_token"), device=torch.device("cuda"), max_batch_size=1, + hf_cache_dir=hf_cache_dir(), ) self.loading = False self.text_encoder = TensorRTCLIPTextModel( @@ -78,18 +74,20 @@ def _encode_prompt( self.pipe._encode_prompt = _encode_prompt def teardown(self): - del self.pipe + if hasattr(self, "pipe"): + del self.pipe torch.cuda.empty_cache() gc.collect() - def generate(self, opts: ImageGenerationOptions) -> ImageGenerationResult: + def generate(self, opts: ImageGenerationOptions): self.wait_loading() if opts.seed is None or opts.seed == -1: opts.seed = random.randrange(0, 4294967294, 1) - e2e_tic = time.perf_counter() results = [] + self.pipe.scheduler = self.get_scheduler(opts.scheduler_id) + for i in range(opts.batch_count): manual_seed = opts.seed + i @@ -98,46 +96,48 @@ def callback( timestep: torch.Tensor, latents: torch.Tensor, ): - timesteps = self.pipe.scheduler.timesteps - include = step % 10 == 0 and len(timesteps) - step >= 10 - if include: - factor = 1.0 / 0.18215 * latents - sample_inp = cuda.DeviceView( - ptr=factor.data_ptr(), shape=factor.shape, dtype=np.float32 - ) - images = self.pipe.run_engine("vae", {"latent": sample_inp})[ - "images" - ] - images = self.pipe.decode_images(images) - images = [ - *images, - *list(chain.from_iterable([x for x, _ in results])), - ] - - async def runner(): - data = DenoiseLatentData( - step=step, - preview={utils.img2b64(x): opts for x in images} - if include - else [], - ) - await sio.emit("denoise_latent", data=data.dict()) - - utils.fire_and_forget(runner)() + queue.put(((opts.steps * i) + step,)) generator = torch.Generator(device=self.pipe.device).manual_seed( manual_seed ) - images = self.pipe( - prompt=opts.prompt, - negative_prompt=opts.negative_prompt, - image_height=opts.image_height, - image_width=opts.image_width, - guidance_scale=opts.scale, - num_inference_steps=opts.steps, - generator=generator, - callback=callback, - ) + + queue = Queue() + done = object() + + def callback( + step: int, + timestep: torch.Tensor, + latents: torch.Tensor, + ): + queue.put(((opts.steps * i) + step, results)) + + def on_done(feature): + queue.put(done) + + with ThreadPoolExecutor() as executer: + feature = executer.submit( + self.pipe, + prompt=[opts.prompt] * opts.batch_size, + negative_prompt=[opts.negative_prompt] * opts.batch_size, + image_height=opts.image_height, + image_width=opts.image_width, + guidance_scale=opts.scale, + num_inference_steps=opts.steps, + generator=generator, + callback=callback, + ) + feature.add_done_callback(on_done) + + while True: + data = queue.get() + if data is done: + break + else: + yield data + + images = feature.result() + results.append( ( images, @@ -146,12 +146,7 @@ async def runner(): ), ) ) - - all_perf_time = time.perf_counter() - e2e_tic - result = ImageGenerationResult(images={}, performance=all_perf_time) - for images, opts in results: for img in images: - result.images[utils.img2b64(img)] = opts save_image(img, opts) - return result + yield results diff --git a/modules/shared.py b/modules/shared.py new file mode 100644 index 00000000..9200b090 --- /dev/null +++ b/modules/shared.py @@ -0,0 +1,11 @@ +import os + +from modules import config + +ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def hf_cache_dir(): + cache_dir = os.path.join(config.get("model_dir"), "diffusers") + os.makedirs(cache_dir, exist_ok=True) + return cache_dir diff --git a/modules/tabs/tensorrt.py b/modules/tabs/tensorrt.py new file mode 100644 index 00000000..f2c8b48c --- /dev/null +++ b/modules/tabs/tensorrt.py @@ -0,0 +1,167 @@ +from typing import * + +import gradio as gr + +from api.models.tensorrt import BuildEngineOptions +from modules import config, model_manager +from modules.acceleration.tensorrt.engine import EngineBuilder +from modules.ui import Tab +from modules.utils import tensorrt_is_available + + +class Txt2Img(Tab): + def title(self): + return "TensorRT" + + def sort(self): + return 2 + + def visible(self): + module, version = tensorrt_is_available() + return module and version and config.get("tensorrt") + + def ui(self, outlet): + with gr.Column(): + with gr.Column(): + gr.Markdown("# Build tensorrt engine") + with gr.Row(): + max_batch_size_number = gr.Number( + label="Max batch size", + value=1, + ) + opt_image_height_slider = gr.Slider( + label="Image height", + minimum=1, + maximum=2048, + step=64, + value=512, + ) + opt_image_width_slider = gr.Slider( + label="Image width", minimum=1, maximum=2048, step=64, value=512 + ) + with gr.Row(): + min_latent_resolution_slider = gr.Slider( + label="Min latent resolution", + minimum=1, + maximum=2048, + step=64, + value=256, + ) + max_latent_resolution_slider = gr.Slider( + label="Max latent resolution", + minimum=1, + maximum=2048, + step=64, + value=1024, + ) + with gr.Row(): + build_enable_refit_checkbox = gr.Checkbox( + label="Enable refit", + value=False, + ) + build_static_batch_checkbox = gr.Checkbox( + label="Static batch", + value=False, + ) + build_dynamic_shape_checkbox = gr.Checkbox( + label="Dynamic shape", + value=True, + ) + build_all_tactics_checkbox = gr.Checkbox( + label="All tactics", + value=False, + ) + build_preview_features_checkbox = gr.Checkbox( + label="Preview features", + value=True, + ) + with gr.Row(): + onnx_opset_slider = gr.Slider( + label="ONNX opset", + minimum=7, + maximum=18, + step=1, + value=16, + ) + with gr.Row(): + force_engine_build_checkbox = gr.Checkbox( + label="Force engine build", + value=False, + ) + force_onnx_export_checkbox = gr.Checkbox( + label="Force ONNX export", + value=False, + ) + force_onnx_optimize_checkbox = gr.Checkbox( + label="Force ONNX optimize", + value=False, + ) + status = gr.Textbox( + label="Status", + value="", + interactive=False, + ) + build_button = gr.Button( + label="Build", + variant="primary", + ) + build_button.click( + fn=self.build_engine, + inputs=[ + max_batch_size_number, + opt_image_height_slider, + opt_image_width_slider, + min_latent_resolution_slider, + max_latent_resolution_slider, + build_enable_refit_checkbox, + build_static_batch_checkbox, + build_dynamic_shape_checkbox, + build_all_tactics_checkbox, + build_preview_features_checkbox, + onnx_opset_slider, + force_engine_build_checkbox, + force_onnx_export_checkbox, + force_onnx_optimize_checkbox, + ], + outputs=[status], + ) + + def build_engine( + self, + max_batch_size: int, + opt_image_height: int, + opt_image_width: int, + min_latent_resolution: int, + max_latent_resolution: int, + build_enable_refit: bool, + build_static_batch: bool, + build_dynamic_shape: bool, + build_all_tactics: bool, + build_preview_features: bool, + onnx_opset: int, + force_engine_build: bool, + force_onnx_export: bool, + force_onnx_optimize: bool, + ): + yield "Building engine..." + model_manager.runner.teardown() + opts = BuildEngineOptions( + max_batch_size=max_batch_size, + opt_image_height=opt_image_height, + opt_image_width=opt_image_width, + min_latent_resolution=min_latent_resolution, + max_latent_resolution=max_latent_resolution, + build_enable_refit=build_enable_refit, + build_static_batch=build_static_batch, + build_dynamic_shape=build_dynamic_shape, + build_all_tactics=build_all_tactics, + build_preview_features=build_preview_features, + onnx_opset=onnx_opset, + force_engine_build=force_engine_build, + force_onnx_export=force_onnx_export, + force_onnx_optimize=force_onnx_optimize, + ) + builder = EngineBuilder(opts) + builder.build() + model_manager.runner.activate() + yield "Engine built" diff --git a/modules/tabs/txt2img.py b/modules/tabs/txt2img.py new file mode 100644 index 00000000..00160c6b --- /dev/null +++ b/modules/tabs/txt2img.py @@ -0,0 +1,85 @@ +import gradio as gr + +from api.models.diffusion import ImageGenerationOptions +from modules import model_manager +from modules.components import image_generation_options +from modules.ui import Tab + + +class Txt2Img(Tab): + def title(self): + return "txt2img" + + def sort(self): + return 1 + + def generate_image( + self, + prompt: str, + negative_prompt: str, + sampler_name: str, + sampling_steps: int, + batch_size: int, + batch_count: int, + cfg_scale: float, + width: int = 512, + height: int = 512, + seed: int = -1, + ): + if model_manager.runner is None: + yield None, "Please select a model." + + yield [], "Generating...", gr.Button.update( + value="Generating...", variant="secondary", interactive=False + ) + + count = 0 + + for data in model_manager.runner.generate( + ImageGenerationOptions( + prompt=prompt, + negative_prompt=negative_prompt, + batch_size=batch_size, + batch_count=batch_count, + scheduler_id=sampler_name, + steps=sampling_steps, + scale=cfg_scale, + image_height=height, + image_width=width, + seed=seed, + ) + ): + if type(data) == tuple: + step, preview = data + progress = step / (batch_count * sampling_steps) + previews = [] + for images, opts in preview: + previews.extend(images) + + if len(previews) == count: + update = gr.Gallery.update() + else: + update = gr.Gallery.update(value=previews) + count = len(previews) + yield update, f"Progress: {progress * 100:.2f}%, Step: {step}", gr.Button.update( + value="Generating...", variant="secondary", interactive=False + ) + else: + image = data + + results = [] + for images, opts in image: + results.extend(images) + + yield results, "Finished", gr.Button.update( + value="Generate", variant="primary", interactive=True + ) + + def ui(self, outlet): + generate_button, prompts, options, outputs = image_generation_options.ui() + + generate_button.click( + fn=self.generate_image, + inputs=[*prompts, *options], + outputs=[*outputs, generate_button], + ) diff --git a/modules/ui.py b/modules/ui.py new file mode 100644 index 00000000..122881f6 --- /dev/null +++ b/modules/ui.py @@ -0,0 +1,147 @@ +import importlib +import os +from typing import * + +import gradio as gr +import gradio.routes + +from modules import shared + +from .components import header +from .shared import ROOT_DIR + + +class Tab: + TABS_DIR = os.path.join(ROOT_DIR, "modules", "tabs") + + def __init__(self, filepath: str) -> None: + self.filepath = filepath + + def sort(self): + return 1 + + def title(self): + return "" + + def ui(self, outlet: Callable): + pass + + def visible(self): + return True + + def __call__(self): + children_dir = self.filepath[:-3] + children = [] + + if os.path.isdir(children_dir): + for file in os.listdir(children_dir): + if not file.endswith(".py"): + continue + module_name = file[:-3] + parent = os.path.relpath(Tab.TABS_DIR, Tab.TABS_DIR).replace("/", ".") + + if parent.startswith("."): + parent = parent[1:] + if parent.endswith("."): + parent = parent[:-1] + + children.append( + importlib.import_module(f"modules.tabs.{parent}.{module_name}") + ) + + children = sorted(children, key=lambda x: x.sort()) + + tabs = [] + + for child in children: + attrs = child.__dict__ + tab = [x for x in attrs.values() if issubclass(x, Tab)] + if len(tab) > 0: + tabs.append(tab[0]) + + def outlet(): + with gr.Tabs(): + for tab in tabs: + if not tab.visible(): + continue + with gr.Tab(tab.title()): + tab() + + return self.ui(outlet) + + +def load_tabs() -> List[Tab]: + tabs = [] + files = os.listdir(os.path.join(ROOT_DIR, "modules", "tabs")) + + for file in files: + if not file.endswith(".py"): + continue + module_name = file[:-3] + module = importlib.import_module(f"modules.tabs.{module_name}") + attrs = module.__dict__ + TabClass = [ + x + for x in attrs.values() + if type(x) == type and issubclass(x, Tab) and not x == Tab + ] + if len(TabClass) > 0: + tabs.append((file, TabClass[0])) + + tabs = sorted([TabClass(file) for file, TabClass in tabs], key=lambda x: x.sort()) + return tabs + + +def webpath(fn): + if fn.startswith(ROOT_DIR): + web_path = os.path.relpath(fn, ROOT_DIR).replace("\\", "/") + else: + web_path = os.path.abspath(fn) + + return f"file={web_path}?{os.path.getmtime(fn)}" + + +def javascript_html(): + script_js = os.path.join(ROOT_DIR, "script.js") + head = f'\n' + + return head + + +def css_html(): + return f'' + + +def create_head(): + head = "" + head += css_html() + head += javascript_html() + + def template_response(*args, **kwargs): + res = shared.gradio_template_response_original(*args, **kwargs) + res.body = res.body.replace(b"", f"{head}".encode("utf8")) + res.init_headers() + return res + + gradio.routes.templates.TemplateResponse = template_response + + +def create_ui(): + block = gr.Blocks() + + with block: + header.ui() + with gr.Tabs(elem_id="radiata-root"): + for tab in load_tabs(): + if not tab.visible(): + continue + with gr.Tab(tab.title()): + tab() + + create_head() + + return block + + +if not hasattr(shared, "gradio_template_response_original"): + shared.gradio_template_response_original = gradio.routes.templates.TemplateResponse diff --git a/modules/utils.py b/modules/utils.py index 04101778..01b42c99 100644 --- a/modules/utils.py +++ b/modules/utils.py @@ -2,7 +2,8 @@ import base64 import importlib import io -from typing import List +from distutils.version import LooseVersion +from typing import * import numpy as np import torch @@ -46,6 +47,12 @@ def is_installed(package: str): return spec is not None +def tensorrt_is_available(): + tensorrt = is_installed("tensorrt") + version = LooseVersion("2") > LooseVersion(torch.__version__) + return tensorrt, version + + def fire_and_forget(f): def wrapped(*args, **kwargs): def runner(): diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 5b3946e2..00000000 --- a/requirements.txt +++ /dev/null @@ -1,26 +0,0 @@ -# API Server -fastapi==0.94.1 -uvicorn[standard]==0.20.0 -pywin32; sys_platform == "win32" -python-socketio>=5 -aiohttp -tqdm -packaging - -# Pre-commit -pre-commit - -# TensorRT Diffusion -accelerate -colored -cuda-python -diffusers==0.14.0 -ftfy -matplotlib -onnx==1.13.1 -onnxruntime==1.14.1 ---extra-index-url https://pypi.ngc.nvidia.com -onnx-graphsurgeon==0.3.26 -polygraphy==0.44.2 -scipy -transformers==4.26.1 \ No newline at end of file diff --git a/requirements/base.txt b/requirements/base.txt new file mode 100644 index 00000000..46665f70 --- /dev/null +++ b/requirements/base.txt @@ -0,0 +1,7 @@ +gradio + +accelerate==0.18.0 +diffusers==0.16.1 +transformers==4.28.1 +pillow +toml \ No newline at end of file diff --git a/requirements/dev.txt b/requirements/dev.txt new file mode 100644 index 00000000..c98f422e --- /dev/null +++ b/requirements/dev.txt @@ -0,0 +1,3 @@ +isort +black +pre-commit \ No newline at end of file diff --git a/requirements/tensorrt.txt b/requirements/tensorrt.txt new file mode 100644 index 00000000..67714f20 --- /dev/null +++ b/requirements/tensorrt.txt @@ -0,0 +1,7 @@ +pywin32; sys_platform == "win32" +cuda-python<12 +onnx==1.12.0 +--extra-index-url https://pypi.ngc.nvidia.com +onnx-graphsurgeon==0.3.25 +onnxruntime==1.13.1 +polygraphy==0.43.1 \ No newline at end of file diff --git a/script.js b/script.js new file mode 100644 index 00000000..60fa6f4a --- /dev/null +++ b/script.js @@ -0,0 +1,30 @@ +window['__RADIATA_CONTEXT'] = { + executedOnLoaded: false +} + +/** + * @returns {ShadowRoot | Document} + */ +function gradioApp() { + const elems = document.getElementsByTagName('gradio-app') + const elem = elems.length == 0 ? document : elems[0] + + if (elem !== document) elem.getElementById = (id) => document.getElementById(id) + return elem.shadowRoot ? elem.shadowRoot : elem +} + +document.addEventListener("DOMContentLoaded", function () { + var mutationObserver = new MutationObserver((m) => { + if (window['__RADIATA_CONTEXT'].executedOnLoaded) return + window['__RADIATA_CONTEXT'].executedOnLoaded = true + const interval = setInterval(() => { + const root = gradioApp().getElementById("radiata-root") + if (root) { + clearInterval(interval) + const button = gradioApp().getElementById("inference-mode-reload-button") + button.click() + } + }, 500) + }) + mutationObserver.observe(gradioApp(), { childList: true, subtree: true }) +}) \ No newline at end of file diff --git a/styles.css b/styles.css new file mode 100644 index 00000000..d8b304ef --- /dev/null +++ b/styles.css @@ -0,0 +1,31 @@ + +/* general gradio fixes */ + +:root, .dark{ + --checkbox-label-gap: 0.25em 0.1em; + --section-header-text-size: 12pt; + --block-background-fill: transparent; +} + + +div.gradio-container{ + max-width: unset !important; +} + +@media screen and (min-width: 2500px) { + .image_generation_gallery { + min-height: 768px; + } +} + +.image_generation_gallery img{ + object-fit: scale-down !important; +} + +.tool-button{ + max-width: 2.2em; + min-width: 2.2em !important; + height: 2.4em; + align-self: center; + border-radius: 0.5em; +} \ No newline at end of file diff --git a/webui.py b/webui.py new file mode 100644 index 00000000..20f35749 --- /dev/null +++ b/webui.py @@ -0,0 +1,32 @@ +import os + +if "--tensorrt" in os.environ.get("COMMANDLINE_ARGS", ""): + import tensorrt as trt + + from lib.tensorrt.utilities import TRT_LOGGER + + print(f"TensorRT version: {trt.__version__}") + trt.init_libnvinfer_plugins(TRT_LOGGER, "") + +from modules import config, model_manager, ui + + +def pre_load(): + config.init() + model_manager.init() + model_manager.set_default_model() + + +def webui(): + pre_load() + app = ui.create_ui() + app.queue(64) + app, local_url, share_url = app.launch( + server_name=config.get("host"), + server_port=config.get("port"), + share=config.get("share"), + ) + + +if __name__ == "__main__": + webui()