Add PyInstaller for apps/ (#907)

Build with pyinstaller.exe .\apps\stable_diffusion\web\shark_sd.spec

normal flow works. exe is missing a few json files

apps/stable_diffusion/scripts/__init__.py

@@ -1 +1 @@
-from .txt2img import txt2img_inf
+from apps.stable_diffusion.scripts.txt2img import txt2img_inf

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,
+)

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,
+)

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

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,
+)

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.

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 = {

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,
+)

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):

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:

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,
+)

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,
 )
 

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
 
 

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,

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.

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()

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,

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",

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)