diff --git a/apps/stable_diffusion/scripts/__init__.py b/apps/stable_diffusion/scripts/__init__.py index b4fcdd64..9ae16cca 100644 --- a/apps/stable_diffusion/scripts/__init__.py +++ b/apps/stable_diffusion/scripts/__init__.py @@ -1 +1 @@ -from .txt2img import txt2img_inf +from apps.stable_diffusion.scripts.txt2img import txt2img_inf diff --git a/apps/stable_diffusion/shark_sd.spec b/apps/stable_diffusion/shark_sd.spec new file mode 100644 index 00000000..967bf7c8 --- /dev/null +++ b/apps/stable_diffusion/shark_sd.spec @@ -0,0 +1,78 @@ +# -*- mode: python ; coding: utf-8 -*- +from PyInstaller.utils.hooks import collect_data_files +from PyInstaller.utils.hooks import copy_metadata + +import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5) + +datas = [] +datas += collect_data_files('torch') +datas += copy_metadata('torch') +datas += copy_metadata('tqdm') +datas += copy_metadata('regex') +datas += copy_metadata('requests') +datas += copy_metadata('packaging') +datas += copy_metadata('filelock') +datas += copy_metadata('numpy') +datas += copy_metadata('tokenizers') +datas += copy_metadata('importlib_metadata') +datas += copy_metadata('torchvision') +datas += copy_metadata('torch-mlir') +datas += copy_metadata('diffusers') +datas += copy_metadata('transformers') +datas += copy_metadata('omegaconf') +datas += copy_metadata('safetensors') +datas += collect_data_files('gradio') +datas += collect_data_files('iree') +datas += collect_data_files('google-cloud-storage') +datas += collect_data_files('shark') +datas += [ + ( 'resources/prompts.json', 'resources' ), + ( 'resources/model_db.json', 'resources' ), + ( 'resources/opt_flags.json', 'resources' ), + ( 'resources/base_model.json', 'resources' ), + ( 'web/logos/*', 'logos' ) + ] + +binaries = [] + +block_cipher = None + + +a = Analysis( + ['web/index.py'], + pathex=['.'], + binaries=binaries, + datas=datas, + hiddenimports=['shark', 'shark.*', 'shark.shark_inference', 'shark_inference', 'iree.tools.core', 'gradio', 'apps'], + hookspath=[], + hooksconfig={}, + runtime_hooks=[], + excludes=[], + win_no_prefer_redirects=False, + win_private_assemblies=False, + cipher=block_cipher, + noarchive=False, +) +pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher) + +exe = EXE( + pyz, + a.scripts, + a.binaries, + a.zipfiles, + a.datas, + [], + name='shark_sd', + debug=False, + bootloader_ignore_signals=False, + strip=False, + upx=True, + upx_exclude=[], + runtime_tmpdir=None, + console=True, + disable_windowed_traceback=False, + argv_emulation=False, + target_arch=None, + codesign_identity=None, + entitlements_file=None, +) diff --git a/apps/stable_diffusion/shark_sd_cli.spec b/apps/stable_diffusion/shark_sd_cli.spec new file mode 100644 index 00000000..fba6a196 --- /dev/null +++ b/apps/stable_diffusion/shark_sd_cli.spec @@ -0,0 +1,84 @@ +# -*- mode: python ; coding: utf-8 -*- +from PyInstaller.utils.hooks import collect_data_files +from PyInstaller.utils.hooks import copy_metadata + +import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5) + +import os + +spec_root = os.path.abspath(SPECPATH) +shark_root = os.path.join(spec_root, "../..") +apps_root = os.path.join(spec_root, "../" +print(spec_root) +print(shark_root) +print(apps_root) + +datas = [] +datas += collect_data_files('torch') +datas += copy_metadata('torch') +datas += copy_metadata('tqdm') +datas += copy_metadata('regex') +datas += copy_metadata('requests') +datas += copy_metadata('packaging') +datas += copy_metadata('filelock') +datas += copy_metadata('numpy') +datas += copy_metadata('tokenizers') +datas += copy_metadata('importlib_metadata') +datas += copy_metadata('torch-mlir') +datas += copy_metadata('diffusers') +datas += copy_metadata('transformers') +datas += copy_metadata('omegaconf') +datas += copy_metadata('safetensors') +datas += collect_data_files('iree') +datas += collect_data_files('google-cloud-storage') +datas += collect_data_files('shark') +datas += collect_data_files('apps') +datas += [ + ( 'resources/prompts.json', 'resources' ), + ( 'resources/model_db.json', 'resources' ), + ( 'resources/opt_flags.json', 'resources' ), + ( 'resources/base_model.json', 'resources' ), + ] + +binaries = [] + +block_cipher = None + +a = Analysis( + ['scripts/txt2img.py'], + pathex=[spec_root, shark_root, apps_root], + binaries=binaries, + datas=datas, + hiddenimports=['apps', 'shark', 'shark.*', 'shark.shark_inference', 'shark_inference', 'iree.tools.core' ], + hookspath=[], + hooksconfig={}, + runtime_hooks=[], + excludes=[], + win_no_prefer_redirects=False, + win_private_assemblies=False, + cipher=block_cipher, + noarchive=False, +) +pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher) + +exe = EXE( + pyz, + a.scripts, + a.binaries, + a.zipfiles, + a.datas, + [], + name='shark_sd_cli', + debug=False, + bootloader_ignore_signals=False, + strip=False, + upx=True, + upx_exclude=[], + runtime_tmpdir=None, + console=True, + disable_windowed_traceback=False, + argv_emulation=False, + target_arch=None, + codesign_identity=None, + entitlements_file=None, +) diff --git a/apps/stable_diffusion/src/__init__.py b/apps/stable_diffusion/src/__init__.py index 5184720b..81afd606 100644 --- a/apps/stable_diffusion/src/__init__.py +++ b/apps/stable_diffusion/src/__init__.py @@ -1,8 +1,8 @@ -from .utils import ( +from apps.stable_diffusion.src.utils import ( args, set_init_device_flags, prompt_examples, get_available_devices, ) -from .pipelines import Text2ImagePipeline -from .schedulers import get_schedulers +from apps.stable_diffusion.src.pipelines import Text2ImagePipeline +from apps.stable_diffusion.src.schedulers import get_schedulers diff --git a/apps/stable_diffusion/src/models/__init__.py b/apps/stable_diffusion/src/models/__init__.py index 2ddbbe97..421897cd 100644 --- a/apps/stable_diffusion/src/models/__init__.py +++ b/apps/stable_diffusion/src/models/__init__.py @@ -1,2 +1,9 @@ -from .model_wrappers import SharkifyStableDiffusionModel -from .opt_params import get_vae, get_unet, get_clip, get_tokenizer +from apps.stable_diffusion.src.models.model_wrappers import ( + SharkifyStableDiffusionModel, +) +from apps.stable_diffusion.src.models.opt_params import ( + get_vae, + get_unet, + get_clip, + get_tokenizer, +) diff --git a/apps/stable_diffusion/src/models/model_wrappers.py b/apps/stable_diffusion/src/models/model_wrappers.py index f3682306..48246485 100644 --- a/apps/stable_diffusion/src/models/model_wrappers.py +++ b/apps/stable_diffusion/src/models/model_wrappers.py @@ -5,7 +5,12 @@ import torch import sys import traceback import re -from ..utils import compile_through_fx, get_opt_flags, base_models, args +from apps.stable_diffusion.src.utils import ( + compile_through_fx, + get_opt_flags, + base_models, + args, +) # These shapes are parameter dependent. diff --git a/apps/stable_diffusion/src/models/opt_params.py b/apps/stable_diffusion/src/models/opt_params.py index caf9a768..ce64601e 100644 --- a/apps/stable_diffusion/src/models/opt_params.py +++ b/apps/stable_diffusion/src/models/opt_params.py @@ -1,6 +1,6 @@ import sys from transformers import CLIPTokenizer -from ..utils import models_db, args, get_shark_model +from apps.stable_diffusion.src.utils import models_db, args, get_shark_model hf_model_variant_map = { diff --git a/apps/stable_diffusion/src/pipelines/__init__.py b/apps/stable_diffusion/src/pipelines/__init__.py index 9ed9f966..75dfe779 100644 --- a/apps/stable_diffusion/src/pipelines/__init__.py +++ b/apps/stable_diffusion/src/pipelines/__init__.py @@ -1 +1,3 @@ -from .pipeline_shark_stable_diffusion_txt2img import Text2ImagePipeline +from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_txt2img import ( + Text2ImagePipeline, +) diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py index 8fc96359..f7ea025d 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py @@ -13,8 +13,10 @@ from diffusers import ( EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, ) -from ..schedulers import SharkEulerDiscreteScheduler -from .pipeline_shark_stable_diffusion_utils import StableDiffusionPipeline +from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler +from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import ( + StableDiffusionPipeline, +) class Text2ImagePipeline(StableDiffusionPipeline): diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py index 45f46ddd..79c95747 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py @@ -13,15 +13,19 @@ from diffusers import ( DPMSolverMultistepScheduler, ) from shark.shark_inference import SharkInference -from ..schedulers import SharkEulerDiscreteScheduler -from ..models import ( +from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler +from apps.stable_diffusion.src.models import ( SharkifyStableDiffusionModel, get_vae, get_clip, get_unet, get_tokenizer, ) -from ..utils import start_profiling, end_profiling, preprocessCKPT +from apps.stable_diffusion.src.utils import ( + start_profiling, + end_profiling, + preprocessCKPT, +) class StableDiffusionPipeline: diff --git a/apps/stable_diffusion/src/schedulers/__init__.py b/apps/stable_diffusion/src/schedulers/__init__.py index d143edc6..8a939b0f 100644 --- a/apps/stable_diffusion/src/schedulers/__init__.py +++ b/apps/stable_diffusion/src/schedulers/__init__.py @@ -1,2 +1,4 @@ -from .sd_schedulers import get_schedulers -from .shark_eulerdiscrete import SharkEulerDiscreteScheduler +from apps.stable_diffusion.src.schedulers.sd_schedulers import get_schedulers +from apps.stable_diffusion.src.schedulers.shark_eulerdiscrete import ( + SharkEulerDiscreteScheduler, +) diff --git a/apps/stable_diffusion/src/schedulers/sd_schedulers.py b/apps/stable_diffusion/src/schedulers/sd_schedulers.py index 514fce4c..8b8809da 100644 --- a/apps/stable_diffusion/src/schedulers/sd_schedulers.py +++ b/apps/stable_diffusion/src/schedulers/sd_schedulers.py @@ -6,7 +6,7 @@ from diffusers import ( EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, ) -from .shark_eulerdiscrete import ( +from apps.stable_diffusion.src.schedulers.shark_eulerdiscrete import ( SharkEulerDiscreteScheduler, ) diff --git a/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py b/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py index 29494d4c..2b8fea65 100644 --- a/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py +++ b/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py @@ -9,7 +9,11 @@ from diffusers import ( EulerDiscreteScheduler, ) from diffusers.configuration_utils import register_to_config -from ..utils import compile_through_fx, get_shark_model, args +from apps.stable_diffusion.src.utils import ( + compile_through_fx, + get_shark_model, + args, +) import torch diff --git a/apps/stable_diffusion/src/utils/__init__.py b/apps/stable_diffusion/src/utils/__init__.py index f2108583..e3aa973f 100644 --- a/apps/stable_diffusion/src/utils/__init__.py +++ b/apps/stable_diffusion/src/utils/__init__.py @@ -1,13 +1,16 @@ -from .profiler import start_profiling, end_profiling -from .resources import ( +from apps.stable_diffusion.src.utils.profiler import ( + start_profiling, + end_profiling, +) +from apps.stable_diffusion.src.utils.resources import ( prompt_examples, models_db, base_models, opt_flags, resource_path, ) -from .stable_args import args -from .utils import ( +from apps.stable_diffusion.src.utils.stable_args import args +from apps.stable_diffusion.src.utils.utils import ( get_shark_model, compile_through_fx, set_iree_runtime_flags, diff --git a/apps/stable_diffusion/src/utils/profiler.py b/apps/stable_diffusion/src/utils/profiler.py index c2045576..1783dc7c 100644 --- a/apps/stable_diffusion/src/utils/profiler.py +++ b/apps/stable_diffusion/src/utils/profiler.py @@ -1,4 +1,4 @@ -from .stable_args import args +from apps.stable_diffusion.src.utils.stable_args import args # Helper function to profile the vulkan device. diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py index 74c868c5..0aa9ff4e 100644 --- a/apps/stable_diffusion/src/utils/stable_args.py +++ b/apps/stable_diffusion/src/utils/stable_args.py @@ -334,4 +334,4 @@ p.add_argument( help="Apply Winograd on selected conv ops.", ) -args = p.parse_args() +args, unknown = p.parse_known_args() diff --git a/apps/stable_diffusion/src/utils/utils.py b/apps/stable_diffusion/src/utils/utils.py index 24b0a282..4a9752a4 100644 --- a/apps/stable_diffusion/src/utils/utils.py +++ b/apps/stable_diffusion/src/utils/utils.py @@ -7,8 +7,8 @@ from shark.iree_utils.vulkan_utils import ( get_vulkan_target_triple, ) from shark.iree_utils.gpu_utils import get_cuda_sm_cc -from .stable_args import args -from .resources import opt_flags +from apps.stable_diffusion.src.utils.stable_args import args +from apps.stable_diffusion.src.utils.resources import opt_flags import sys from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( load_pipeline_from_original_stable_diffusion_ckpt, diff --git a/setup.py b/setup.py index 032842fd..35594bdb 100644 --- a/setup.py +++ b/setup.py @@ -2,11 +2,12 @@ from setuptools import find_packages from setuptools import setup import os +import glob with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() -PACKAGE_VERSION = os.environ.get("SHARK_PACKAGE_VERSION") or "0.0.4" +PACKAGE_VERSION = os.environ.get("SHARK_PACKAGE_VERSION") or "0.0.5" backend_deps = [] if "NO_BACKEND" in os.environ.keys(): backend_deps = [ @@ -34,6 +35,7 @@ setup( ], packages=find_packages(exclude=("examples")), python_requires=">=3.9", + data_files=glob.glob("apps/stable_diffusion/resources/**"), install_requires=[ "numpy", "PyYAML", diff --git a/web/models/stable_diffusion/main.py b/web/models/stable_diffusion/main.py index 94803cee..2f72dd92 100644 --- a/web/models/stable_diffusion/main.py +++ b/web/models/stable_diffusion/main.py @@ -226,7 +226,6 @@ def stable_diff_inf( avg_ms = 0 for i, t in tqdm(enumerate(scheduler.timesteps)): - step_start = time.time() timestep = torch.tensor([t]).to(dtype).detach().numpy() latent_model_input = scheduler.scale_model_input(latents, t)