[Web] Upgrade Gradio to v3.21.0 (#1188)

Co-authored-by: powderluv <powderluv@users.noreply.github.com>

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+count = 1
+show-source = 1
+select = E9,F63,F7,F82
+exclude = lit.cfg.py

.github/workflows/test-models.yml

@@ -99,11 +99,12 @@ jobs:
       run: |
         # black format check
         black --version
-        black --line-length 79 --check .
+        black --check .
         # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude lit.cfg.py
+        flake8 . --statistics
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude lit.cfg.py
+        flake8 . --isolated --count --exit-zero --max-complexity=10 --max-line-length=127 \
+          --statistics --exclude lit.cfg.py
 
     - name: Validate Models on CPU
       if: matrix.suite == 'cpu'

.gitignore

@@ -168,6 +168,8 @@ shark_tmp/
 *.vmfb
 .use-iree
 tank/dict_configs.py
+*.csv
+reproducers/
 
 # ORT related artefacts
 cache_models/
@@ -182,3 +184,6 @@ models/
 
 # models folder
 apps/stable_diffusion/web/models/
+
+# Stencil annotators.
+stencil_annotator/

.style.yapf

@@ -1,3 +0,0 @@
-[style]
-  based_on_style = google
-  column_limit = 80

apps/stable_diffusion/scripts/img2img.py

@@ -38,11 +38,38 @@ init_use_tuned = args.use_tuned
 init_import_mlir = args.import_mlir
 
 
+# For stencil, the input image can be of any size but we need to ensure that
+# it conforms with our model contraints :-
+#   Both width and height should be > 384 and multiple of 8.
+# This utility function performs the transformation on the input image while
+# also maintaining the aspect ratio before sending it to the stencil pipeline.
+def resize_stencil(image: Image.Image):
+    width, height = image.size
+    aspect_ratio = width / height
+    min_size = min(width, height)
+    if min_size < 384:
+        n_size = 384
+        if width == min_size:
+            width = n_size
+            height = n_size / aspect_ratio
+        else:
+            height = n_size
+            width = n_size * aspect_ratio
+    width = int(width)
+    height = int(height)
+    n_width = width // 8
+    n_height = height // 8
+    n_width *= 8
+    n_height *= 8
+    new_image = image.resize((n_width, n_height))
+    return new_image, n_width, n_height
+
+
 # Exposed to UI.
 def img2img_inf(
     prompt: str,
     negative_prompt: str,
-    init_image: Image,
+    init_image,
     height: int,
     width: int,
     steps: int,
@@ -61,6 +88,8 @@ def img2img_inf(
     save_metadata_to_json: bool,
     save_metadata_to_png: bool,
 ):
+    from apps.stable_diffusion.web.ui.utils import get_custom_model_pathfile
+
     global img2img_obj
     global config_obj
     global schedulers
@@ -93,7 +122,7 @@ def img2img_inf(
             )
         args.hf_model_id = hf_model_id
     elif ".ckpt" in custom_model or ".safetensors" in custom_model:
-        args.ckpt_loc = custom_model
+        args.ckpt_loc = get_custom_model_pathfile(custom_model)
     else:
         args.hf_model_id = custom_model
 
@@ -105,6 +134,7 @@ def img2img_inf(
     if use_stencil is not None:
         args.scheduler = "DDIM"
         args.hf_model_id = "runwayml/stable-diffusion-v1-5"
+        image, width, height = resize_stencil(image)
     elif args.scheduler != "PNDM":
         if "Shark" in args.scheduler:
             print(
@@ -147,6 +177,7 @@ def img2img_inf(
         )
         schedulers = get_schedulers(model_id)
         scheduler_obj = schedulers[scheduler]
+
         if use_stencil is not None:
             args.use_tuned = False
             img2img_obj = StencilPipeline.from_pretrained(
@@ -164,6 +195,7 @@ def img2img_inf(
                 args.use_tuned,
                 low_cpu_mem_usage=args.low_cpu_mem_usage,
                 use_stencil=use_stencil,
+                debug=args.import_debug if args.import_mlir else False,
             )
         else:
             img2img_obj = Image2ImagePipeline.from_pretrained(
@@ -180,6 +212,7 @@ def img2img_inf(
                 args.use_base_vae,
                 args.use_tuned,
                 low_cpu_mem_usage=args.low_cpu_mem_usage,
+                debug=args.import_debug if args.import_mlir else False,
             )
 
     img2img_obj.scheduler = schedulers[scheduler]
@@ -214,6 +247,7 @@ def img2img_inf(
         generated_imgs.extend(out_imgs)
         seeds.append(img_seed)
         img2img_obj.log += "\n"
+        yield generated_imgs, generated_imgs[0], img2img_obj.log
 
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"
@@ -225,7 +259,7 @@ def img2img_inf(
     text_output += img2img_obj.log
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-    return generated_imgs, text_output
+    yield generated_imgs, text_output
 
 
 if __name__ == "__main__":
@@ -236,6 +270,7 @@ if __name__ == "__main__":
         print("Flag --img_path is required.")
         exit()
 
+    image = Image.open(args.img_path).convert("RGB")
     # When the models get uploaded, it should be default to False.
     args.import_mlir = True
 
@@ -243,6 +278,7 @@ if __name__ == "__main__":
     if use_stencil:
         args.scheduler = "DDIM"
         args.hf_model_id = "runwayml/stable-diffusion-v1-5"
+        image, args.width, args.height = resize_stencil(image)
     elif args.scheduler != "PNDM":
         if "Shark" in args.scheduler:
             print(
@@ -257,9 +293,7 @@ if __name__ == "__main__":
     dtype = torch.float32 if args.precision == "fp32" else torch.half
     set_init_device_flags()
     schedulers = get_schedulers(args.hf_model_id)
-
     scheduler_obj = schedulers[args.scheduler]
-    image = Image.open(args.img_path).convert("RGB")
     seed = utils.sanitize_seed(args.seed)
     # Adjust for height and width based on model
 
@@ -279,6 +313,7 @@ if __name__ == "__main__":
             args.use_tuned,
             low_cpu_mem_usage=args.low_cpu_mem_usage,
             use_stencil=use_stencil,
+            debug=args.import_debug if args.import_mlir else False,
         )
     else:
         img2img_obj = Image2ImagePipeline.from_pretrained(
@@ -295,6 +330,7 @@ if __name__ == "__main__":
             args.use_base_vae,
             args.use_tuned,
             low_cpu_mem_usage=args.low_cpu_mem_usage,
+            debug=args.import_debug if args.import_mlir else False,
         )
 
     start_time = time.time()

apps/stable_diffusion/scripts/inpaint.py

@@ -1,4 +1,3 @@
-import sys
 import torch
 import time
 from PIL import Image
@@ -30,6 +29,11 @@ inpaint_obj = None
 config_obj = None
 schedulers = None
 
+# set initial values of iree_vulkan_target_triple, use_tuned and import_mlir.
+init_iree_vulkan_target_triple = args.iree_vulkan_target_triple
+init_use_tuned = args.use_tuned
+init_import_mlir = args.import_mlir
+
 
 # Exposed to UI.
 def inpaint_inf(
@@ -38,6 +42,8 @@ def inpaint_inf(
     image_dict,
     height: int,
     width: int,
+    inpaint_full_res: bool,
+    inpaint_full_res_padding: int,
     steps: int,
     guidance_scale: float,
     seed: int,
@@ -52,6 +58,8 @@ def inpaint_inf(
     save_metadata_to_json: bool,
     save_metadata_to_png: bool,
 ):
+    from apps.stable_diffusion.web.ui.utils import get_custom_model_pathfile
+
     global inpaint_obj
     global config_obj
     global schedulers
@@ -79,7 +87,7 @@ def inpaint_inf(
             )
         args.hf_model_id = hf_model_id
     elif ".ckpt" in custom_model or ".safetensors" in custom_model:
-        args.ckpt_loc = custom_model
+        args.ckpt_loc = get_custom_model_pathfile(custom_model)
     else:
         args.hf_model_id = custom_model
 
@@ -106,9 +114,9 @@ def inpaint_inf(
         args.height = height
         args.width = width
         args.device = device.split("=>", 1)[1].strip()
-        args.iree_vulkan_target_triple = ""
-        args.use_tuned = True
-        args.import_mlir = False
+        args.iree_vulkan_target_triple = init_iree_vulkan_target_triple
+        args.use_tuned = init_use_tuned
+        args.import_mlir = init_import_mlir
         set_init_device_flags()
         model_id = (
             args.hf_model_id
@@ -118,18 +126,20 @@ def inpaint_inf(
         schedulers = get_schedulers(model_id)
         scheduler_obj = schedulers[scheduler]
         inpaint_obj = InpaintPipeline.from_pretrained(
-            scheduler_obj,
-            args.import_mlir,
-            args.hf_model_id,
-            args.ckpt_loc,
-            args.custom_vae,
-            args.precision,
-            args.max_length,
-            args.batch_size,
-            args.height,
-            args.width,
-            args.use_base_vae,
-            args.use_tuned,
+            scheduler=scheduler_obj,
+            import_mlir=args.import_mlir,
+            model_id=args.hf_model_id,
+            ckpt_loc=args.ckpt_loc,
+            precision=args.precision,
+            max_length=args.max_length,
+            batch_size=args.batch_size,
+            height=args.height,
+            width=args.width,
+            use_base_vae=args.use_base_vae,
+            use_tuned=args.use_tuned,
+            custom_vae=args.custom_vae,
+            low_cpu_mem_usage=args.low_cpu_mem_usage,
+            debug=args.import_debug if args.import_mlir else False,
         )
 
     inpaint_obj.scheduler = schedulers[scheduler]
@@ -152,6 +162,8 @@ def inpaint_inf(
             batch_size,
             height,
             width,
+            inpaint_full_res,
+            inpaint_full_res_padding,
             steps,
             guidance_scale,
             img_seed,
@@ -164,6 +176,7 @@ def inpaint_inf(
         generated_imgs.extend(out_imgs)
         seeds.append(img_seed)
         inpaint_obj.log += "\n"
+        yield generated_imgs, generated_imgs[0], inpaint_obj.log
 
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"
@@ -175,7 +188,7 @@ def inpaint_inf(
     text_output += inpaint_obj.log
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-    return generated_imgs, text_output
+    yield generated_imgs, text_output
 
 
 if __name__ == "__main__":
@@ -204,18 +217,20 @@ if __name__ == "__main__":
     mask_image = Image.open(args.mask_path)
 
     inpaint_obj = InpaintPipeline.from_pretrained(
-        scheduler_obj,
-        args.import_mlir,
-        args.hf_model_id,
-        args.ckpt_loc,
-        args.custom_vae,
-        args.precision,
-        args.max_length,
-        args.batch_size,
-        args.height,
-        args.width,
-        args.use_base_vae,
-        args.use_tuned,
+        scheduler=scheduler_obj,
+        import_mlir=args.import_mlir,
+        model_id=args.hf_model_id,
+        ckpt_loc=args.ckpt_loc,
+        precision=args.precision,
+        max_length=args.max_length,
+        batch_size=args.batch_size,
+        height=args.height,
+        width=args.width,
+        use_base_vae=args.use_base_vae,
+        use_tuned=args.use_tuned,
+        custom_vae=args.custom_vae,
+        low_cpu_mem_usage=args.low_cpu_mem_usage,
+        debug=args.import_debug if args.import_mlir else False,
     )
 
     for current_batch in range(args.batch_count):
@@ -232,6 +247,8 @@ if __name__ == "__main__":
             args.batch_size,
             args.height,
             args.width,
+            args.inpaint_full_res,
+            args.inpaint_full_res_padding,
             args.steps,
             args.guidance_scale,
             seed,

apps/stable_diffusion/scripts/outpaint.py

@@ -1,4 +1,3 @@
-import sys
 import torch
 import time
 from PIL import Image
@@ -40,7 +39,7 @@ init_import_mlir = args.import_mlir
 def outpaint_inf(
     prompt: str,
     negative_prompt: str,
-    init_image: Image,
+    init_image,
     pixels: int,
     mask_blur: int,
     directions: list,
@@ -62,6 +61,8 @@ def outpaint_inf(
     save_metadata_to_json: bool,
     save_metadata_to_png: bool,
 ):
+    from apps.stable_diffusion.web.ui.utils import get_custom_model_pathfile
+
     global outpaint_obj
     global config_obj
     global schedulers
@@ -88,7 +89,7 @@ def outpaint_inf(
             )
         args.hf_model_id = hf_model_id
     elif ".ckpt" in custom_model or ".safetensors" in custom_model:
-        args.ckpt_loc = custom_model
+        args.ckpt_loc = get_custom_model_pathfile(custom_model)
     else:
         args.hf_model_id = custom_model
 
@@ -184,6 +185,7 @@ def outpaint_inf(
         generated_imgs.extend(out_imgs)
         seeds.append(img_seed)
         outpaint_obj.log += "\n"
+        yield generated_imgs, generated_imgs[0], outpaint_obj.log
 
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"
@@ -195,7 +197,7 @@ def outpaint_inf(
     text_output += outpaint_obj.log
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-    return generated_imgs, text_output
+    yield generated_imgs, text_output
 
 
 if __name__ == "__main__":

apps/stable_diffusion/scripts/txt2img.py

@@ -1,4 +1,3 @@
-import sys
 import torch
 import time
 from dataclasses import dataclass
@@ -23,6 +22,7 @@ class Config:
     height: int
     width: int
     device: str
+    use_lora: str
 
 
 txt2img_obj = None
@@ -54,7 +54,11 @@ def txt2img_inf(
     max_length: int,
     save_metadata_to_json: bool,
     save_metadata_to_png: bool,
+    lora_weights: str,
+    lora_hf_id: str,
 ):
+    from apps.stable_diffusion.web.ui.utils import get_custom_model_pathfile
+
     global txt2img_obj
     global config_obj
     global schedulers
@@ -80,13 +84,22 @@ def txt2img_inf(
             )
         args.hf_model_id = hf_model_id
     elif ".ckpt" in custom_model or ".safetensors" in custom_model:
-        args.ckpt_loc = custom_model
+        args.ckpt_loc = get_custom_model_pathfile(custom_model)
     else:
         args.hf_model_id = custom_model
 
     args.save_metadata_to_json = save_metadata_to_json
     args.write_metadata_to_png = save_metadata_to_png
 
+    use_lora = ""
+    if lora_weights == "None" and not lora_hf_id:
+        use_lora = ""
+    elif not lora_hf_id:
+        use_lora = lora_weights
+    else:
+        use_lora = lora_hf_id
+    args.use_lora = use_lora
+
     dtype = torch.float32 if precision == "fp32" else torch.half
     cpu_scheduling = not scheduler.startswith("Shark")
     new_config_obj = Config(
@@ -98,6 +111,7 @@ def txt2img_inf(
         height,
         width,
         device,
+        use_lora,
     )
     if not txt2img_obj or config_obj != new_config_obj:
         config_obj = new_config_obj
@@ -120,19 +134,21 @@ def txt2img_inf(
         schedulers = get_schedulers(model_id)
         scheduler_obj = schedulers[scheduler]
         txt2img_obj = Text2ImagePipeline.from_pretrained(
-            scheduler_obj,
-            args.import_mlir,
-            args.hf_model_id,
-            args.ckpt_loc,
-            args.custom_vae,
-            args.precision,
-            args.max_length,
-            args.batch_size,
-            args.height,
-            args.width,
-            args.use_base_vae,
-            args.use_tuned,
+            scheduler=scheduler_obj,
+            import_mlir=args.import_mlir,
+            model_id=args.hf_model_id,
+            ckpt_loc=args.ckpt_loc,
+            precision=args.precision,
+            max_length=args.max_length,
+            batch_size=args.batch_size,
+            height=args.height,
+            width=args.width,
+            use_base_vae=args.use_base_vae,
+            use_tuned=args.use_tuned,
+            custom_vae=args.custom_vae,
             low_cpu_mem_usage=args.low_cpu_mem_usage,
+            debug=args.import_debug if args.import_mlir else False,
+            use_lora=use_lora,
         )
 
     txt2img_obj.scheduler = schedulers[scheduler]
@@ -190,21 +206,23 @@ if __name__ == "__main__":
     schedulers = get_schedulers(args.hf_model_id)
     scheduler_obj = schedulers[args.scheduler]
     seed = args.seed
-
+    use_lora = args.use_lora
     txt2img_obj = Text2ImagePipeline.from_pretrained(
-        scheduler_obj,
-        args.import_mlir,
-        args.hf_model_id,
-        args.ckpt_loc,
-        args.custom_vae,
-        args.precision,
-        args.max_length,
-        args.batch_size,
-        args.height,
-        args.width,
-        args.use_base_vae,
-        args.use_tuned,
+        scheduler=scheduler_obj,
+        import_mlir=args.import_mlir,
+        model_id=args.hf_model_id,
+        ckpt_loc=args.ckpt_loc,
+        precision=args.precision,
+        max_length=args.max_length,
+        batch_size=args.batch_size,
+        height=args.height,
+        width=args.width,
+        use_base_vae=args.use_base_vae,
+        use_tuned=args.use_tuned,
+        custom_vae=args.custom_vae,
         low_cpu_mem_usage=args.low_cpu_mem_usage,
+        debug=args.import_debug if args.import_mlir else False,
+        use_lora=use_lora,
     )
 
     for current_batch in range(args.batch_count):

apps/stable_diffusion/shark_sd.spec

@@ -1,6 +1,7 @@
 # -*- mode: python ; coding: utf-8 -*-
 from PyInstaller.utils.hooks import collect_data_files
 from PyInstaller.utils.hooks import copy_metadata
+from PyInstaller.utils.hooks import collect_submodules
 
 import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5)
 
@@ -21,6 +22,7 @@ datas += copy_metadata('safetensors')
 datas += collect_data_files('diffusers')
 datas += collect_data_files('transformers')
 datas += collect_data_files('opencv-python')
+datas += collect_data_files('skimage')
 datas += collect_data_files('gradio')
 datas += collect_data_files('iree')
 datas += collect_data_files('google-cloud-storage')
@@ -38,13 +40,15 @@ binaries = []
 
 block_cipher = None
 
+hiddenimports = ['shark', 'shark.shark_inference', 'apps']
+hiddenimports += [x for x in collect_submodules("skimage") if "tests" not in x]
 
 a = Analysis(
     ['web/index.py'],
     pathex=['.'],
     binaries=binaries,
     datas=datas,
-    hiddenimports=['shark', 'shark.shark_inference', 'apps'],
+    hiddenimports=hiddenimports,
     hookspath=[],
     hooksconfig={},
     runtime_hooks=[],

apps/stable_diffusion/shark_sd_cli.spec

@@ -1,5 +1,6 @@
 # -*- mode: python ; coding: utf-8 -*-
 from PyInstaller.utils.hooks import collect_data_files
+from PyInstaller.utils.hooks import collect_submodules
 from PyInstaller.utils.hooks import copy_metadata
 
 import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5)
@@ -21,6 +22,7 @@ datas += copy_metadata('safetensors')
 datas += collect_data_files('diffusers')
 datas += collect_data_files('transformers')
 datas += collect_data_files('opencv-python')
+datas += collect_data_files('skimage')
 datas += collect_data_files('gradio')
 datas += collect_data_files('iree')
 datas += collect_data_files('google-cloud-storage')
@@ -36,13 +38,15 @@ binaries = []
 
 block_cipher = None
 
+hiddenimports = ['shark', 'shark.shark_inference', 'apps']
+hiddenimports += [x for x in collect_submodules("skimage") if "tests" not in x]
 
 a = Analysis(
     ['scripts/txt2img.py'],
     pathex=['.'],
     binaries=binaries,
     datas=datas,
-    hiddenimports=['shark', 'shark.shark_inference', 'apps'],
+    hiddenimports=hiddenimports,
     hookspath=[],
     hooksconfig={},
     runtime_hooks=[],

apps/stable_diffusion/src/models/model_wrappers.py

@@ -5,6 +5,7 @@ import torch
 import safetensors.torch
 import traceback
 import sys
+import os
 from apps.stable_diffusion.src.utils import (
     compile_through_fx,
     get_opt_flags,
@@ -16,6 +17,7 @@ from apps.stable_diffusion.src.utils import (
     fetch_and_update_base_model_id,
     get_path_stem,
     get_extended_name,
+    get_stencil_model_id,
 )
 
 
@@ -30,13 +32,23 @@ def replace_shape_str(shape, max_len, width, height, batch_size):
         elif shape[i] == "width":
             new_shape.append(width)
         elif isinstance(shape[i], str):
-            mul_val = int(shape[i].split("*")[0])
-            if "batch_size" in shape[i]:
-                new_shape.append(batch_size * mul_val)
-            elif "height" in shape[i]:
-                new_shape.append(height * mul_val)
-            elif "width" in shape[i]:
-                new_shape.append(width * mul_val)
+            if "*" in shape[i]:
+                mul_val = int(shape[i].split("*")[0])
+                if "batch_size" in shape[i]:
+                    new_shape.append(batch_size * mul_val)
+                elif "height" in shape[i]:
+                    new_shape.append(height * mul_val)
+                elif "width" in shape[i]:
+                    new_shape.append(width * mul_val)
+            elif "/" in shape[i]:
+                import math
+                div_val = int(shape[i].split("/")[1])
+                if "batch_size" in shape[i]:
+                    new_shape.append(math.ceil(batch_size / div_val))
+                elif "height" in shape[i]:
+                    new_shape.append(math.ceil(height / div_val))
+                elif "width" in shape[i]:
+                    new_shape.append(math.ceil(width / div_val))
         else:
             new_shape.append(shape[i])
     return new_shape
@@ -81,7 +93,12 @@ class SharkifyStableDiffusionModel:
         use_base_vae: bool = False,
         use_tuned: bool = False,
         low_cpu_mem_usage: bool = False,
-        is_inpaint: bool = False
+        debug: bool = False,
+        sharktank_dir: str = "",
+        generate_vmfb: bool = True,
+        is_inpaint: bool = False,
+        use_stencil: str = None,
+        use_lora: str = ""
     ):
         self.check_params(max_len, width, height)
         self.max_len = max_len
@@ -102,7 +119,8 @@ class SharkifyStableDiffusionModel:
         self.precision = precision
         self.base_vae = use_base_vae
         self.model_name = (
-            str(batch_size)
+            "_"
+            + str(batch_size)
             + "_"
             + str(max_len)
             + "_"
@@ -112,12 +130,22 @@ class SharkifyStableDiffusionModel:
             + "_"
             + precision
         )
+        print(f'use_tuned? sharkify: {use_tuned}')
         self.use_tuned = use_tuned
         if use_tuned:
             self.model_name = self.model_name + "_tuned"
         self.model_name = self.model_name + "_" + get_path_stem(self.model_id)
         self.low_cpu_mem_usage = low_cpu_mem_usage
         self.is_inpaint = is_inpaint
+        self.use_stencil = get_stencil_model_id(use_stencil)
+        if use_lora != "":
+            self.model_name = self.model_name + "_" + get_path_stem(use_lora)
+        self.use_lora = use_lora
+
+        print(self.model_name)
+        self.debug = debug
+        self.sharktank_dir = sharktank_dir
+        self.generate_vmfb = generate_vmfb
 
     def get_extended_name_for_all_model(self, mask_to_fetch):
         model_name = {}
@@ -212,12 +240,18 @@ class SharkifyStableDiffusionModel:
         vae = VaeModel(low_cpu_mem_usage=self.low_cpu_mem_usage)
         inputs = tuple(self.inputs["vae"])
         is_f16 = True if self.precision == "fp16" else False
+        save_dir = os.path.join(self.sharktank_dir, self.model_name["vae"])
+        if self.debug:
+            os.makedirs(save_dir, exist_ok=True)
         shark_vae = compile_through_fx(
             vae,
             inputs,
             is_f16=is_f16,
             use_tuned=self.use_tuned,
             model_name=self.model_name["vae"],
+            debug=self.debug,
+            generate_vmfb=self.generate_vmfb,
+            save_dir=save_dir,
             extra_args=get_opt_flags("vae", precision=self.precision),
         )
         return shark_vae
@@ -229,7 +263,7 @@ class SharkifyStableDiffusionModel:
             ):
                 super().__init__()
                 self.unet = UNet2DConditionModel.from_pretrained(
-                    "takuma104/control_sd15_canny",  # TODO: ADD with model ID
+                    model_id,
                     subfolder="unet",
                     low_cpu_mem_usage=low_cpu_mem_usage,
                 )
@@ -277,12 +311,11 @@ class SharkifyStableDiffusionModel:
     def get_control_net(self):
         class StencilControlNetModel(torch.nn.Module):
             def __init__(
-                self, model_id=self.model_id, low_cpu_mem_usage=False
+                self, model_id=self.use_stencil, low_cpu_mem_usage=False
             ):
                 super().__init__()
                 self.cnet = ControlNetModel.from_pretrained(
-                    "takuma104/control_sd15_canny", # TODO: ADD with model ID
-                    subfolder="controlnet",
+                    model_id,
                     low_cpu_mem_usage=low_cpu_mem_usage,
                 )
                 self.in_channels = self.cnet.in_channels
@@ -330,13 +363,15 @@ class SharkifyStableDiffusionModel:
 
     def get_unet(self):
         class UnetModel(torch.nn.Module):
-            def __init__(self, model_id=self.model_id, low_cpu_mem_usage=False):
+            def __init__(self, model_id=self.model_id, low_cpu_mem_usage=False, use_lora=self.use_lora):
                 super().__init__()
                 self.unet = UNet2DConditionModel.from_pretrained(
                     model_id,
                     subfolder="unet",
                     low_cpu_mem_usage=low_cpu_mem_usage,
                 )
+                if use_lora != "":
+                    self.unet.load_attn_procs(use_lora)
                 self.in_channels = self.unet.in_channels
                 self.train(False)
                 if(args.attention_slicing is not None and args.attention_slicing != "none"):
@@ -364,6 +399,12 @@ class SharkifyStableDiffusionModel:
         is_f16 = True if self.precision == "fp16" else False
         inputs = tuple(self.inputs["unet"])
         input_mask = [True, True, True, False]
+        save_dir = os.path.join(self.sharktank_dir, self.model_name["unet"])
+        if self.debug:
+            os.makedirs(
+                save_dir,
+                exist_ok=True,
+            )
         shark_unet = compile_through_fx(
             unet,
             inputs,
@@ -371,6 +412,9 @@ class SharkifyStableDiffusionModel:
             is_f16=is_f16,
             f16_input_mask=input_mask,
             use_tuned=self.use_tuned,
+            debug=self.debug,
+            generate_vmfb=self.generate_vmfb,
+            save_dir=save_dir,
             extra_args=get_opt_flags("unet", precision=self.precision),
         )
         return shark_unet
@@ -389,10 +433,19 @@ class SharkifyStableDiffusionModel:
                 return self.text_encoder(input)[0]
 
         clip_model = CLIPText(low_cpu_mem_usage=self.low_cpu_mem_usage)
+        save_dir = os.path.join(self.sharktank_dir, self.model_name["clip"])
+        if self.debug:
+            os.makedirs(
+                save_dir,
+                exist_ok=True,
+            )
         shark_clip = compile_through_fx(
             clip_model,
             tuple(self.inputs["clip"]),
             model_name=self.model_name["clip"],
+            debug=self.debug,
+            generate_vmfb=self.generate_vmfb,
+            save_dir=save_dir,
             extra_args=get_opt_flags("clip", precision="fp32"),
         )
         return shark_clip
@@ -454,7 +507,7 @@ class SharkifyStableDiffusionModel:
         # --  Fetch all vmfbs for the model, if present, else delete the lot.
         need_vae_encode, need_stencil = False, False
         if args.img_path is not None:
-            if args.use_stencil is not None:
+            if self.use_stencil is not None:
                 need_stencil = True
             else:
                 need_vae_encode = True

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py

@@ -2,7 +2,7 @@ import torch
 from tqdm.auto import tqdm
 import numpy as np
 from random import randint
-from PIL import Image
+from PIL import Image, ImageOps
 from transformers import CLIPTokenizer
 from typing import Union
 from shark.shark_inference import SharkInference
@@ -43,10 +43,222 @@ class InpaintPipeline(StableDiffusionPipeline):
         super().__init__(vae, text_encoder, tokenizer, unet, scheduler)
         self.vae_encode = vae_encode
 
-    def prepare_mask_and_masked_image(self, image, mask, height, width):
+    def prepare_latents(
+        self,
+        batch_size,
+        height,
+        width,
+        generator,
+        num_inference_steps,
+        dtype,
+    ):
+        latents = torch.randn(
+            (
+                batch_size,
+                4,
+                height // 8,
+                width // 8,
+            ),
+            generator=generator,
+            dtype=torch.float32,
+        ).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def get_crop_region(self, mask, pad=0):
+        h, w = mask.shape
+
+        crop_left = 0
+        for i in range(w):
+            if not (mask[:, i] == 0).all():
+                break
+            crop_left += 1
+
+        crop_right = 0
+        for i in reversed(range(w)):
+            if not (mask[:, i] == 0).all():
+                break
+            crop_right += 1
+
+        crop_top = 0
+        for i in range(h):
+            if not (mask[i] == 0).all():
+                break
+            crop_top += 1
+
+        crop_bottom = 0
+        for i in reversed(range(h)):
+            if not (mask[i] == 0).all():
+                break
+            crop_bottom += 1
+
+        return (
+            int(max(crop_left - pad, 0)),
+            int(max(crop_top - pad, 0)),
+            int(min(w - crop_right + pad, w)),
+            int(min(h - crop_bottom + pad, h)),
+        )
+
+    def expand_crop_region(
+        self,
+        crop_region,
+        processing_width,
+        processing_height,
+        image_width,
+        image_height,
+    ):
+        x1, y1, x2, y2 = crop_region
+
+        ratio_crop_region = (x2 - x1) / (y2 - y1)
+        ratio_processing = processing_width / processing_height
+
+        if ratio_crop_region > ratio_processing:
+            desired_height = (x2 - x1) / ratio_processing
+            desired_height_diff = int(desired_height - (y2 - y1))
+            y1 -= desired_height_diff // 2
+            y2 += desired_height_diff - desired_height_diff // 2
+            if y2 >= image_height:
+                diff = y2 - image_height
+                y2 -= diff
+                y1 -= diff
+            if y1 < 0:
+                y2 -= y1
+                y1 -= y1
+            if y2 >= image_height:
+                y2 = image_height
+        else:
+            desired_width = (y2 - y1) * ratio_processing
+            desired_width_diff = int(desired_width - (x2 - x1))
+            x1 -= desired_width_diff // 2
+            x2 += desired_width_diff - desired_width_diff // 2
+            if x2 >= image_width:
+                diff = x2 - image_width
+                x2 -= diff
+                x1 -= diff
+            if x1 < 0:
+                x2 -= x1
+                x1 -= x1
+            if x2 >= image_width:
+                x2 = image_width
+
+        return x1, y1, x2, y2
+
+    def resize_image(self, resize_mode, im, width, height):
+        """
+        resize_mode:
+            0: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.
+            1: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.
+        """
+
+        if resize_mode == 0:
+            ratio = width / height
+            src_ratio = im.width / im.height
+
+            src_w = (
+                width if ratio > src_ratio else im.width * height // im.height
+            )
+            src_h = (
+                height if ratio <= src_ratio else im.height * width // im.width
+            )
+
+            resized = im.resize((src_w, src_h), resample=Image.LANCZOS)
+            res = Image.new("RGB", (width, height))
+            res.paste(
+                resized,
+                box=(width // 2 - src_w // 2, height // 2 - src_h // 2),
+            )
+
+        else:
+            ratio = width / height
+            src_ratio = im.width / im.height
+
+            src_w = (
+                width if ratio < src_ratio else im.width * height // im.height
+            )
+            src_h = (
+                height if ratio >= src_ratio else im.height * width // im.width
+            )
+
+            resized = im.resize((src_w, src_h), resample=Image.LANCZOS)
+            res = Image.new("RGB", (width, height))
+            res.paste(
+                resized,
+                box=(width // 2 - src_w // 2, height // 2 - src_h // 2),
+            )
+
+            if ratio < src_ratio:
+                fill_height = height // 2 - src_h // 2
+                res.paste(
+                    resized.resize((width, fill_height), box=(0, 0, width, 0)),
+                    box=(0, 0),
+                )
+                res.paste(
+                    resized.resize(
+                        (width, fill_height),
+                        box=(0, resized.height, width, resized.height),
+                    ),
+                    box=(0, fill_height + src_h),
+                )
+            elif ratio > src_ratio:
+                fill_width = width // 2 - src_w // 2
+                res.paste(
+                    resized.resize(
+                        (fill_width, height), box=(0, 0, 0, height)
+                    ),
+                    box=(0, 0),
+                )
+                res.paste(
+                    resized.resize(
+                        (fill_width, height),
+                        box=(resized.width, 0, resized.width, height),
+                    ),
+                    box=(fill_width + src_w, 0),
+                )
+
+        return res
+
+    def prepare_mask_and_masked_image(
+        self,
+        image,
+        mask,
+        height,
+        width,
+        inpaint_full_res,
+        inpaint_full_res_padding,
+    ):
         # preprocess image
         image = image.resize((width, height))
         mask = mask.resize((width, height))
+
+        paste_to = ()
+        overlay_image = None
+        if inpaint_full_res:
+            # prepare overlay image
+            overlay_image = Image.new("RGB", (image.width, image.height))
+            overlay_image.paste(
+                image.convert("RGB"),
+                mask=ImageOps.invert(mask.convert("L")),
+            )
+
+            # prepare mask
+            mask = mask.convert("L")
+            crop_region = self.get_crop_region(
+                np.array(mask), inpaint_full_res_padding
+            )
+            crop_region = self.expand_crop_region(
+                crop_region, width, height, mask.width, mask.height
+            )
+            x1, y1, x2, y2 = crop_region
+            mask = mask.crop(crop_region)
+            mask = self.resize_image(1, mask, width, height)
+            paste_to = (x1, y1, x2 - x1, y2 - y1)
+
+            # prepare image
+            image = image.crop(crop_region)
+            image = self.resize_image(1, image, width, height)
+
         if isinstance(image, (Image.Image, np.ndarray)):
             image = [image]
 
@@ -77,32 +289,7 @@ class InpaintPipeline(StableDiffusionPipeline):
 
         masked_image = image * (mask < 0.5)
 
-        return mask, masked_image
-
-    def prepare_latents(
-        self,
-        batch_size,
-        height,
-        width,
-        generator,
-        num_inference_steps,
-        dtype,
-    ):
-        latents = torch.randn(
-            (
-                batch_size,
-                4,
-                height // 8,
-                width // 8,
-            ),
-            generator=generator,
-            dtype=torch.float32,
-        ).to(dtype)
-
-        self.scheduler.set_timesteps(num_inference_steps)
-        self.scheduler.is_scale_input_called = True
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
+        return mask, masked_image, paste_to, overlay_image
 
     def prepare_mask_latents(
         self,
@@ -143,6 +330,13 @@ class InpaintPipeline(StableDiffusionPipeline):
             )
         return mask, masked_image_latents
 
+    def apply_overlay(self, image, paste_loc, overlay):
+        x, y, w, h = paste_loc
+        image = self.resize_image(0, image, w, h)
+        overlay.paste(image, (x, y))
+
+        return overlay
+
     def generate_images(
         self,
         prompts,
@@ -152,6 +346,8 @@ class InpaintPipeline(StableDiffusionPipeline):
         batch_size,
         height,
         width,
+        inpaint_full_res,
+        inpaint_full_res_padding,
         num_inference_steps,
         guidance_scale,
         seed,
@@ -194,8 +390,18 @@ class InpaintPipeline(StableDiffusionPipeline):
         guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
 
         # Preprocess mask and image
-        mask, masked_image = self.prepare_mask_and_masked_image(
-            image, mask_image, height, width
+        (
+            mask,
+            masked_image,
+            paste_to,
+            overlay_image,
+        ) = self.prepare_mask_and_masked_image(
+            image,
+            mask_image,
+            height,
+            width,
+            inpaint_full_res,
+            inpaint_full_res_padding,
         )
 
         # Prepare mask latent variables
@@ -230,4 +436,10 @@ class InpaintPipeline(StableDiffusionPipeline):
             )
             all_imgs.extend(imgs)
 
+        if inpaint_full_res:
+            output_image = self.apply_overlay(
+                all_imgs[0], paste_to, overlay_image
+            )
+            return [output_image]
+
         return all_imgs

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py

@@ -65,7 +65,6 @@ class OutpaintPipeline(StableDiffusionPipeline):
         ).to(dtype)
 
         self.scheduler.set_timesteps(num_inference_steps)
-        self.scheduler.is_scale_input_called = True
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py

@@ -20,9 +20,6 @@ from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils i
     StableDiffusionPipeline,
 )
 
-import cv2
-from PIL import Image
-
 
 class Text2ImagePipeline(StableDiffusionPipeline):
     def __init__(

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py

@@ -161,15 +161,6 @@ class StableDiffusionPipeline:
                 ),
                 send_to_host=False,
             )
-            down_block_res_samples = control[0:12]
-            mid_block_res_sample = control[12:]
-            down_block_res_samples = [
-                down_block_res_sample * controlnet_conditioning_scale
-                for down_block_res_sample in down_block_res_samples
-            ]
-            mid_block_res_sample = (
-                mid_block_res_sample[0] * controlnet_conditioning_scale
-            )
             timestep = timestep.detach().numpy()
             # Profiling Unet.
             profile_device = start_profiling(file_path="unet.rdc")
@@ -181,19 +172,19 @@ class StableDiffusionPipeline:
                     timestep,
                     text_embeddings_numpy,
                     guidance_scale,
-                    down_block_res_samples[0],
-                    down_block_res_samples[1],
-                    down_block_res_samples[2],
-                    down_block_res_samples[3],
-                    down_block_res_samples[4],
-                    down_block_res_samples[5],
-                    down_block_res_samples[6],
-                    down_block_res_samples[7],
-                    down_block_res_samples[8],
-                    down_block_res_samples[9],
-                    down_block_res_samples[10],
-                    down_block_res_samples[11],
-                    mid_block_res_sample,
+                    control[0],
+                    control[1],
+                    control[2],
+                    control[3],
+                    control[4],
+                    control[5],
+                    control[6],
+                    control[7],
+                    control[8],
+                    control[9],
+                    control[10],
+                    control[11],
+                    control[12],
                 ),
                 send_to_host=False,
             )
@@ -317,7 +308,9 @@ class StableDiffusionPipeline:
         use_base_vae: bool,
         use_tuned: bool,
         low_cpu_mem_usage: bool = False,
-        use_stencil: bool = False,
+        debug: bool = False,
+        use_stencil: str = None,
+        use_lora: str = "",
     ):
         is_inpaint = cls.__name__ in [
             "InpaintPipeline",
@@ -336,7 +329,10 @@ class StableDiffusionPipeline:
                 use_base_vae=use_base_vae,
                 use_tuned=use_tuned,
                 low_cpu_mem_usage=low_cpu_mem_usage,
+                debug=debug,
                 is_inpaint=is_inpaint,
+                use_stencil=use_stencil,
+                use_lora=use_lora,
             )
             if cls.__name__ in [
                 "Image2ImagePipeline",

apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py

@@ -90,8 +90,8 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
         def _import(self):
             scaling_model = ScalingModel()
             self.scaling_model = compile_through_fx(
-                scaling_model,
-                (example_latent, example_sigma),
+                model=scaling_model,
+                inputs=(example_latent, example_sigma),
                 model_name=f"euler_scale_model_input_{BATCH_SIZE}_{args.height}_{args.width}"
                 + args.precision,
                 extra_args=iree_flags,

apps/stable_diffusion/src/utils/__init__.py

@@ -13,6 +13,7 @@ from apps.stable_diffusion.src.utils.sd_annotation import sd_model_annotation
 from apps.stable_diffusion.src.utils.stable_args import args
 from apps.stable_diffusion.src.utils.stencils.stencil_utils import (
     controlnet_hint_conversion,
+    get_stencil_model_id,
 )
 from apps.stable_diffusion.src.utils.utils import (
     get_shark_model,

apps/stable_diffusion/src/utils/resources/base_model.json

@@ -110,7 +110,7 @@
                 "dtype": "f32"
             },
             "controlnet_hint": {
-                "shape": [1, 3, 512, 512],
+                "shape": [1, 3, "8*height", "8*width"],
                 "dtype": "f32"
             }
         },
@@ -143,55 +143,55 @@
                 "dtype": "f32"
             },
             "control1": {
-                "shape": [2, 320, 64, 64],
+                "shape": [2, 320, "height", "width"],
                 "dtype": "f32"
             },
             "control2": {
-                "shape": [2, 320, 64, 64],
+                "shape": [2, 320, "height", "width"],
                 "dtype": "f32"
             },
             "control3": {
-                "shape": [2, 320, 64, 64],
+                "shape": [2, 320, "height", "width"],
                 "dtype": "f32"
             },
             "control4": {
-                "shape": [2, 320, 32, 32],
+                "shape": [2, 320, "height/2", "width/2"],
                 "dtype": "f32"
             },
             "control5": {
-                "shape": [2, 640, 32, 32],
+                "shape": [2, 640, "height/2", "width/2"],
                 "dtype": "f32"
             },
             "control6": {
-                "shape": [2, 640, 32, 32],
+                "shape": [2, 640, "height/2", "width/2"],
                 "dtype": "f32"
             },
             "control7": {
-                "shape": [2, 640, 16, 16],
+                "shape": [2, 640, "height/4", "width/4"],
                 "dtype": "f32"
             },
             "control8": {
-                "shape": [2, 1280, 16, 16],
+                "shape": [2, 1280, "height/4", "width/4"],
                 "dtype": "f32"
             },
             "control9": {
-                "shape": [2, 1280, 16, 16],
+                "shape": [2, 1280, "height/4", "width/4"],
                 "dtype": "f32"
             },
             "control10": {
-                "shape": [2, 1280, 8, 8],
+                "shape": [2, 1280, "height/8", "width/8"],
                 "dtype": "f32"
             },
             "control11": {
-                "shape": [2, 1280, 8, 8],
+                "shape": [2, 1280, "height/8", "width/8"],
                 "dtype": "f32"
             },
             "control12": {
-                "shape": [2, 1280, 8, 8],
+                "shape": [2, 1280, "height/8", "width/8"],
                 "dtype": "f32"
             },
             "control13": {
-                "shape": [2, 1280, 8, 8],
+                "shape": [2, 1280, "height/8", "width/8"],
                 "dtype": "f32"
             }
         },

apps/stable_diffusion/src/utils/resources/model_db.json

@@ -18,10 +18,15 @@
     "stablediffusion/v1_4/unet/fp16/length_77/tuned":"unet_8dec_fp16_tuned",
     "stablediffusion/v1_4/unet/fp16/length_77/tuned/cuda":"unet_8dec_fp16_cuda_tuned",
     "stablediffusion/v1_4/unet/fp32/length_77/untuned":"unet_1dec_fp32",
+    "stablediffusion/v1_4/unet/fp32/length_64/untuned":"unet_1_64_512_512_fp32_CompVis_stable_diffusion_v1_4",
     "stablediffusion/v1_4/vae/fp16/length_77/untuned":"vae_19dec_fp16",
     "stablediffusion/v1_4/vae/fp16/length_77/tuned":"vae_19dec_fp16_tuned",
     "stablediffusion/v1_4/vae/fp16/length_77/tuned/cuda":"vae_19dec_fp16_cuda_tuned",
     "stablediffusion/v1_4/vae/fp16/length_77/untuned/base":"vae_8dec_fp16",
+    "stablediffusion/v1_4/vae/fp32/length_77/untuned":"vae_1_64_512_512_fp32_CompVis_stable_diffusion_v1_4",
+    "stablediffusion/v1_4/vae/fp32/length_64/untuned":"vae_1_64_512_512_fp32_CompVis_stable_diffusion_v1_4",
+    "stablediffusion/v1_4/clip/fp32/length_77/untuned":"clip_18dec_fp32",
+    "stablediffusion/v1_4/clip/fp32/length_64/untuned":"clip_1_64_512_512_fp32_CompVis_stable_diffusion_v1_4",
     "stablediffusion/v2_1base/unet/fp16/length_77/untuned":"unet77_512_512_fp16_stabilityai_stable_diffusion_2_1_base",
     "stablediffusion/v2_1base/unet/fp16/length_77/tuned":"unet2base_8dec_fp16_tuned_v2",
     "stablediffusion/v2_1base/unet/fp16/length_77/tuned/cuda":"unet2base_8dec_fp16_cuda_tuned",

apps/stable_diffusion/src/utils/stable_args.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 from pathlib import Path
 
 
@@ -6,6 +7,13 @@ def path_expand(s):
     return Path(s).expanduser().resolve()
 
 
+def is_valid_file(arg):
+    if not os.path.exists(arg):
+        return None
+    else:
+        return arg
+
+
 p = argparse.ArgumentParser(
     description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
 )
@@ -104,6 +112,21 @@ p.add_argument(
     help="Path to the mask image input for inpainting",
 )
 
+p.add_argument(
+    "--inpaint_full_res",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="If inpaint only masked area or whole picture",
+)
+
+p.add_argument(
+    "--inpaint_full_res_padding",
+    type=int,
+    default=32,
+    choices=range(0, 257, 4),
+    help="Number of pixels for only masked padding",
+)
+
 p.add_argument(
     "--pixels",
     type=int,
@@ -274,10 +297,17 @@ p.add_argument(
 
 p.add_argument(
     "--use_stencil",
-    choices=["canny"],
+    choices=["canny", "openpose", "scribble"],
     help="Enable the stencil feature.",
 )
 
+p.add_argument(
+    "--use_lora",
+    type=str,
+    default="",
+    help="Use standalone LoRA weight using a HF ID or a checkpoint file (~3 MB)",
+)
+
 ##############################################################################
 ### IREE - Vulkan supported flags
 ##############################################################################
@@ -387,6 +417,12 @@ p.add_argument(
     help="flag for whether or not to save generation information in PNG chunk text to generated images.",
 )
 
+p.add_argument(
+    "--import_debug",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="if import_mlir is True, saves mlir via the debug option in shark importer. Does nothing if import_mlir is false (the default)",
+)
 ##############################################################################
 ### Web UI flags
 ##############################################################################
@@ -446,3 +482,7 @@ p.add_argument(
 )
 
 args, unknown = p.parse_known_args()
+if args.import_debug:
+    os.environ["IREE_SAVE_TEMPS"] = os.path.join(
+        os.getcwd(), args.hf_model_id.replace("/", "_")
+    )

apps/stable_diffusion/src/utils/stencils/__init__.py

@@ -0,0 +1,2 @@
+from apps.stable_diffusion.src.utils.stencils.canny import CannyDetector
+from apps.stable_diffusion.src.utils.stencils.openpose import OpenposeDetector

apps/stable_diffusion/src/utils/stencils/openpose/__init__.py

@@ -0,0 +1,62 @@
+import requests
+from pathlib import Path
+
+import torch
+import numpy as np
+
+# from annotator.util import annotator_ckpts_path
+from apps.stable_diffusion.src.utils.stencils.openpose.body import Body
+from apps.stable_diffusion.src.utils.stencils.openpose.hand import Hand
+from apps.stable_diffusion.src.utils.stencils.openpose.openpose_util import (
+    draw_bodypose,
+    draw_handpose,
+    handDetect,
+)
+
+
+body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth"
+hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth"
+
+
+class OpenposeDetector:
+    def __init__(self):
+        cwd = Path.cwd()
+        ckpt_path = Path(cwd, "stencil_annotator")
+        ckpt_path.mkdir(parents=True, exist_ok=True)
+        body_modelpath = ckpt_path / "body_pose_model.pth"
+        hand_modelpath = ckpt_path / "hand_pose_model.pth"
+
+        if not body_modelpath.is_file():
+            r = requests.get(body_model_path, allow_redirects=True)
+            open(body_modelpath, "wb").write(r.content)
+        if not hand_modelpath.is_file():
+            r = requests.get(hand_model_path, allow_redirects=True)
+            open(hand_modelpath, "wb").write(r.content)
+
+        self.body_estimation = Body(body_modelpath)
+        self.hand_estimation = Hand(hand_modelpath)
+
+    def __call__(self, oriImg, hand=False):
+        oriImg = oriImg[:, :, ::-1].copy()
+        with torch.no_grad():
+            candidate, subset = self.body_estimation(oriImg)
+            canvas = np.zeros_like(oriImg)
+            canvas = draw_bodypose(canvas, candidate, subset)
+            if hand:
+                hands_list = handDetect(candidate, subset, oriImg)
+                all_hand_peaks = []
+                for x, y, w, is_left in hands_list:
+                    peaks = self.hand_estimation(
+                        oriImg[y : y + w, x : x + w, :]
+                    )
+                    peaks[:, 0] = np.where(
+                        peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x
+                    )
+                    peaks[:, 1] = np.where(
+                        peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y
+                    )
+                    all_hand_peaks.append(peaks)
+                canvas = draw_handpose(canvas, all_hand_peaks)
+            return canvas, dict(
+                candidate=candidate.tolist(), subset=subset.tolist()
+            )

apps/stable_diffusion/src/utils/stencils/openpose/body.py

@@ -0,0 +1,499 @@
+import cv2
+import numpy as np
+import math
+from scipy.ndimage.filters import gaussian_filter
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+from apps.stable_diffusion.src.utils.stencils.openpose.openpose_util import (
+    make_layers,
+    transfer,
+    padRightDownCorner,
+)
+
+
+class BodyPoseModel(nn.Module):
+    def __init__(self):
+        super(BodyPoseModel, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = [
+            "conv5_5_CPM_L1",
+            "conv5_5_CPM_L2",
+            "Mconv7_stage2_L1",
+            "Mconv7_stage2_L2",
+            "Mconv7_stage3_L1",
+            "Mconv7_stage3_L2",
+            "Mconv7_stage4_L1",
+            "Mconv7_stage4_L2",
+            "Mconv7_stage5_L1",
+            "Mconv7_stage5_L2",
+            "Mconv7_stage6_L1",
+            "Mconv7_stage6_L1",
+        ]
+        blocks = {}
+        block0 = OrderedDict(
+            [
+                ("conv1_1", [3, 64, 3, 1, 1]),
+                ("conv1_2", [64, 64, 3, 1, 1]),
+                ("pool1_stage1", [2, 2, 0]),
+                ("conv2_1", [64, 128, 3, 1, 1]),
+                ("conv2_2", [128, 128, 3, 1, 1]),
+                ("pool2_stage1", [2, 2, 0]),
+                ("conv3_1", [128, 256, 3, 1, 1]),
+                ("conv3_2", [256, 256, 3, 1, 1]),
+                ("conv3_3", [256, 256, 3, 1, 1]),
+                ("conv3_4", [256, 256, 3, 1, 1]),
+                ("pool3_stage1", [2, 2, 0]),
+                ("conv4_1", [256, 512, 3, 1, 1]),
+                ("conv4_2", [512, 512, 3, 1, 1]),
+                ("conv4_3_CPM", [512, 256, 3, 1, 1]),
+                ("conv4_4_CPM", [256, 128, 3, 1, 1]),
+            ]
+        )
+
+        # Stage 1
+        block1_1 = OrderedDict(
+            [
+                ("conv5_1_CPM_L1", [128, 128, 3, 1, 1]),
+                ("conv5_2_CPM_L1", [128, 128, 3, 1, 1]),
+                ("conv5_3_CPM_L1", [128, 128, 3, 1, 1]),
+                ("conv5_4_CPM_L1", [128, 512, 1, 1, 0]),
+                ("conv5_5_CPM_L1", [512, 38, 1, 1, 0]),
+            ]
+        )
+
+        block1_2 = OrderedDict(
+            [
+                ("conv5_1_CPM_L2", [128, 128, 3, 1, 1]),
+                ("conv5_2_CPM_L2", [128, 128, 3, 1, 1]),
+                ("conv5_3_CPM_L2", [128, 128, 3, 1, 1]),
+                ("conv5_4_CPM_L2", [128, 512, 1, 1, 0]),
+                ("conv5_5_CPM_L2", [512, 19, 1, 1, 0]),
+            ]
+        )
+        blocks["block1_1"] = block1_1
+        blocks["block1_2"] = block1_2
+
+        self.model0 = make_layers(block0, no_relu_layers)
+
+        # Stages 2 - 6
+        for i in range(2, 7):
+            blocks["block%d_1" % i] = OrderedDict(
+                [
+                    ("Mconv1_stage%d_L1" % i, [185, 128, 7, 1, 3]),
+                    ("Mconv2_stage%d_L1" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv3_stage%d_L1" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv4_stage%d_L1" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv5_stage%d_L1" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv6_stage%d_L1" % i, [128, 128, 1, 1, 0]),
+                    ("Mconv7_stage%d_L1" % i, [128, 38, 1, 1, 0]),
+                ]
+            )
+
+            blocks["block%d_2" % i] = OrderedDict(
+                [
+                    ("Mconv1_stage%d_L2" % i, [185, 128, 7, 1, 3]),
+                    ("Mconv2_stage%d_L2" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv3_stage%d_L2" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv4_stage%d_L2" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv5_stage%d_L2" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv6_stage%d_L2" % i, [128, 128, 1, 1, 0]),
+                    ("Mconv7_stage%d_L2" % i, [128, 19, 1, 1, 0]),
+                ]
+            )
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_1 = blocks["block1_1"]
+        self.model2_1 = blocks["block2_1"]
+        self.model3_1 = blocks["block3_1"]
+        self.model4_1 = blocks["block4_1"]
+        self.model5_1 = blocks["block5_1"]
+        self.model6_1 = blocks["block6_1"]
+
+        self.model1_2 = blocks["block1_2"]
+        self.model2_2 = blocks["block2_2"]
+        self.model3_2 = blocks["block3_2"]
+        self.model4_2 = blocks["block4_2"]
+        self.model5_2 = blocks["block5_2"]
+        self.model6_2 = blocks["block6_2"]
+
+    def forward(self, x):
+        out1 = self.model0(x)
+
+        out1_1 = self.model1_1(out1)
+        out1_2 = self.model1_2(out1)
+        out2 = torch.cat([out1_1, out1_2, out1], 1)
+
+        out2_1 = self.model2_1(out2)
+        out2_2 = self.model2_2(out2)
+        out3 = torch.cat([out2_1, out2_2, out1], 1)
+
+        out3_1 = self.model3_1(out3)
+        out3_2 = self.model3_2(out3)
+        out4 = torch.cat([out3_1, out3_2, out1], 1)
+
+        out4_1 = self.model4_1(out4)
+        out4_2 = self.model4_2(out4)
+        out5 = torch.cat([out4_1, out4_2, out1], 1)
+
+        out5_1 = self.model5_1(out5)
+        out5_2 = self.model5_2(out5)
+        out6 = torch.cat([out5_1, out5_2, out1], 1)
+
+        out6_1 = self.model6_1(out6)
+        out6_2 = self.model6_2(out6)
+
+        return out6_1, out6_2
+
+
+class Body(object):
+    def __init__(self, model_path):
+        self.model = BodyPoseModel()
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+        model_dict = transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre1 = 0.1
+        thre2 = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
+        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = cv2.resize(
+                oriImg,
+                (0, 0),
+                fx=scale,
+                fy=scale,
+                interpolation=cv2.INTER_CUBIC,
+            )
+            imageToTest_padded, pad = padRightDownCorner(
+                imageToTest, stride, padValue
+            )
+            im = (
+                np.transpose(
+                    np.float32(imageToTest_padded[:, :, :, np.newaxis]),
+                    (3, 2, 0, 1),
+                )
+                / 256
+                - 0.5
+            )
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            with torch.no_grad():
+                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
+            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
+            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            heatmap = np.transpose(
+                np.squeeze(Mconv7_stage6_L2), (1, 2, 0)
+            )  # output 1 is heatmaps
+            heatmap = cv2.resize(
+                heatmap,
+                (0, 0),
+                fx=stride,
+                fy=stride,
+                interpolation=cv2.INTER_CUBIC,
+            )
+            heatmap = heatmap[
+                : imageToTest_padded.shape[0] - pad[2],
+                : imageToTest_padded.shape[1] - pad[3],
+                :,
+            ]
+            heatmap = cv2.resize(
+                heatmap,
+                (oriImg.shape[1], oriImg.shape[0]),
+                interpolation=cv2.INTER_CUBIC,
+            )
+
+            # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
+            paf = np.transpose(
+                np.squeeze(Mconv7_stage6_L1), (1, 2, 0)
+            )  # output 0 is PAFs
+            paf = cv2.resize(
+                paf,
+                (0, 0),
+                fx=stride,
+                fy=stride,
+                interpolation=cv2.INTER_CUBIC,
+            )
+            paf = paf[
+                : imageToTest_padded.shape[0] - pad[2],
+                : imageToTest_padded.shape[1] - pad[3],
+                :,
+            ]
+            paf = cv2.resize(
+                paf,
+                (oriImg.shape[1], oriImg.shape[0]),
+                interpolation=cv2.INTER_CUBIC,
+            )
+
+            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
+            paf_avg += +paf / len(multiplier)
+
+        all_peaks = []
+        peak_counter = 0
+
+        for part in range(18):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+
+            map_left = np.zeros(one_heatmap.shape)
+            map_left[1:, :] = one_heatmap[:-1, :]
+            map_right = np.zeros(one_heatmap.shape)
+            map_right[:-1, :] = one_heatmap[1:, :]
+            map_up = np.zeros(one_heatmap.shape)
+            map_up[:, 1:] = one_heatmap[:, :-1]
+            map_down = np.zeros(one_heatmap.shape)
+            map_down[:, :-1] = one_heatmap[:, 1:]
+
+            peaks_binary = np.logical_and.reduce(
+                (
+                    one_heatmap >= map_left,
+                    one_heatmap >= map_right,
+                    one_heatmap >= map_up,
+                    one_heatmap >= map_down,
+                    one_heatmap > thre1,
+                )
+            )
+            peaks = list(
+                zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])
+            )  # note reverse
+            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
+            peak_id = range(peak_counter, peak_counter + len(peaks))
+            peaks_with_score_and_id = [
+                peaks_with_score[i] + (peak_id[i],)
+                for i in range(len(peak_id))
+            ]
+
+            all_peaks.append(peaks_with_score_and_id)
+            peak_counter += len(peaks)
+
+        # find connection in the specified sequence, center 29 is in the position 15
+        limbSeq = [
+            [2, 3],
+            [2, 6],
+            [3, 4],
+            [4, 5],
+            [6, 7],
+            [7, 8],
+            [2, 9],
+            [9, 10],
+            [10, 11],
+            [2, 12],
+            [12, 13],
+            [13, 14],
+            [2, 1],
+            [1, 15],
+            [15, 17],
+            [1, 16],
+            [16, 18],
+            [3, 17],
+            [6, 18],
+        ]
+        # the middle joints heatmap correpondence
+        mapIdx = [
+            [31, 32],
+            [39, 40],
+            [33, 34],
+            [35, 36],
+            [41, 42],
+            [43, 44],
+            [19, 20],
+            [21, 22],
+            [23, 24],
+            [25, 26],
+            [27, 28],
+            [29, 30],
+            [47, 48],
+            [49, 50],
+            [53, 54],
+            [51, 52],
+            [55, 56],
+            [37, 38],
+            [45, 46],
+        ]
+
+        connection_all = []
+        special_k = []
+        mid_num = 10
+
+        for k in range(len(mapIdx)):
+            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
+            candA = all_peaks[limbSeq[k][0] - 1]
+            candB = all_peaks[limbSeq[k][1] - 1]
+            nA = len(candA)
+            nB = len(candB)
+            indexA, indexB = limbSeq[k]
+            if nA != 0 and nB != 0:
+                connection_candidate = []
+                for i in range(nA):
+                    for j in range(nB):
+                        vec = np.subtract(candB[j][:2], candA[i][:2])
+                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
+                        norm = max(0.001, norm)
+                        vec = np.divide(vec, norm)
+
+                        startend = list(
+                            zip(
+                                np.linspace(
+                                    candA[i][0], candB[j][0], num=mid_num
+                                ),
+                                np.linspace(
+                                    candA[i][1], candB[j][1], num=mid_num
+                                ),
+                            )
+                        )
+
+                        vec_x = np.array(
+                            [
+                                score_mid[
+                                    int(round(startend[I][1])),
+                                    int(round(startend[I][0])),
+                                    0,
+                                ]
+                                for I in range(len(startend))
+                            ]
+                        )
+                        vec_y = np.array(
+                            [
+                                score_mid[
+                                    int(round(startend[I][1])),
+                                    int(round(startend[I][0])),
+                                    1,
+                                ]
+                                for I in range(len(startend))
+                            ]
+                        )
+
+                        score_midpts = np.multiply(
+                            vec_x, vec[0]
+                        ) + np.multiply(vec_y, vec[1])
+                        score_with_dist_prior = sum(score_midpts) / len(
+                            score_midpts
+                        ) + min(0.5 * oriImg.shape[0] / norm - 1, 0)
+                        criterion1 = len(
+                            np.nonzero(score_midpts > thre2)[0]
+                        ) > 0.8 * len(score_midpts)
+                        criterion2 = score_with_dist_prior > 0
+                        if criterion1 and criterion2:
+                            connection_candidate.append(
+                                [
+                                    i,
+                                    j,
+                                    score_with_dist_prior,
+                                    score_with_dist_prior
+                                    + candA[i][2]
+                                    + candB[j][2],
+                                ]
+                            )
+
+                connection_candidate = sorted(
+                    connection_candidate, key=lambda x: x[2], reverse=True
+                )
+                connection = np.zeros((0, 5))
+                for c in range(len(connection_candidate)):
+                    i, j, s = connection_candidate[c][0:3]
+                    if i not in connection[:, 3] and j not in connection[:, 4]:
+                        connection = np.vstack(
+                            [connection, [candA[i][3], candB[j][3], s, i, j]]
+                        )
+                        if len(connection) >= min(nA, nB):
+                            break
+
+                connection_all.append(connection)
+            else:
+                special_k.append(k)
+                connection_all.append([])
+
+        # last number in each row is the total parts number of that person
+        # the second last number in each row is the score of the overall configuration
+        subset = -1 * np.ones((0, 20))
+        candidate = np.array(
+            [item for sublist in all_peaks for item in sublist]
+        )
+
+        for k in range(len(mapIdx)):
+            if k not in special_k:
+                partAs = connection_all[k][:, 0]
+                partBs = connection_all[k][:, 1]
+                indexA, indexB = np.array(limbSeq[k]) - 1
+
+                for i in range(len(connection_all[k])):  # = 1:size(temp,1)
+                    found = 0
+                    subset_idx = [-1, -1]
+                    for j in range(len(subset)):  # 1:size(subset,1):
+                        if (
+                            subset[j][indexA] == partAs[i]
+                            or subset[j][indexB] == partBs[i]
+                        ):
+                            subset_idx[found] = j
+                            found += 1
+
+                    if found == 1:
+                        j = subset_idx[0]
+                        if subset[j][indexB] != partBs[i]:
+                            subset[j][indexB] = partBs[i]
+                            subset[j][-1] += 1
+                            subset[j][-2] += (
+                                candidate[partBs[i].astype(int), 2]
+                                + connection_all[k][i][2]
+                            )
+                    elif found == 2:  # if found 2 and disjoint, merge them
+                        j1, j2 = subset_idx
+                        membership = (
+                            (subset[j1] >= 0).astype(int)
+                            + (subset[j2] >= 0).astype(int)
+                        )[:-2]
+                        if len(np.nonzero(membership == 2)[0]) == 0:  # merge
+                            subset[j1][:-2] += subset[j2][:-2] + 1
+                            subset[j1][-2:] += subset[j2][-2:]
+                            subset[j1][-2] += connection_all[k][i][2]
+                            subset = np.delete(subset, j2, 0)
+                        else:  # as like found == 1
+                            subset[j1][indexB] = partBs[i]
+                            subset[j1][-1] += 1
+                            subset[j1][-2] += (
+                                candidate[partBs[i].astype(int), 2]
+                                + connection_all[k][i][2]
+                            )
+
+                    # if find no partA in the subset, create a new subset
+                    elif not found and k < 17:
+                        row = -1 * np.ones(20)
+                        row[indexA] = partAs[i]
+                        row[indexB] = partBs[i]
+                        row[-1] = 2
+                        row[-2] = (
+                            sum(
+                                candidate[
+                                    connection_all[k][i, :2].astype(int), 2
+                                ]
+                            )
+                            + connection_all[k][i][2]
+                        )
+                        subset = np.vstack([subset, row])
+        # delete some rows of subset which has few parts occur
+        deleteIdx = []
+        for i in range(len(subset)):
+            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
+                deleteIdx.append(i)
+        subset = np.delete(subset, deleteIdx, axis=0)
+
+        # candidate: x, y, score, id
+        return candidate, subset

apps/stable_diffusion/src/utils/stencils/openpose/hand.py

@@ -0,0 +1,205 @@
+import cv2
+import numpy as np
+from scipy.ndimage.filters import gaussian_filter
+import torch
+import torch.nn as nn
+from skimage.measure import label
+from collections import OrderedDict
+from apps.stable_diffusion.src.utils.stencils.openpose.openpose_util import (
+    make_layers,
+    transfer,
+    padRightDownCorner,
+    npmax,
+)
+
+
+class HandPoseModel(nn.Module):
+    def __init__(self):
+        super(HandPoseModel, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = [
+            "conv6_2_CPM",
+            "Mconv7_stage2",
+            "Mconv7_stage3",
+            "Mconv7_stage4",
+            "Mconv7_stage5",
+            "Mconv7_stage6",
+        ]
+        # stage 1
+        block1_0 = OrderedDict(
+            [
+                ("conv1_1", [3, 64, 3, 1, 1]),
+                ("conv1_2", [64, 64, 3, 1, 1]),
+                ("pool1_stage1", [2, 2, 0]),
+                ("conv2_1", [64, 128, 3, 1, 1]),
+                ("conv2_2", [128, 128, 3, 1, 1]),
+                ("pool2_stage1", [2, 2, 0]),
+                ("conv3_1", [128, 256, 3, 1, 1]),
+                ("conv3_2", [256, 256, 3, 1, 1]),
+                ("conv3_3", [256, 256, 3, 1, 1]),
+                ("conv3_4", [256, 256, 3, 1, 1]),
+                ("pool3_stage1", [2, 2, 0]),
+                ("conv4_1", [256, 512, 3, 1, 1]),
+                ("conv4_2", [512, 512, 3, 1, 1]),
+                ("conv4_3", [512, 512, 3, 1, 1]),
+                ("conv4_4", [512, 512, 3, 1, 1]),
+                ("conv5_1", [512, 512, 3, 1, 1]),
+                ("conv5_2", [512, 512, 3, 1, 1]),
+                ("conv5_3_CPM", [512, 128, 3, 1, 1]),
+            ]
+        )
+
+        block1_1 = OrderedDict(
+            [
+                ("conv6_1_CPM", [128, 512, 1, 1, 0]),
+                ("conv6_2_CPM", [512, 22, 1, 1, 0]),
+            ]
+        )
+
+        blocks = {}
+        blocks["block1_0"] = block1_0
+        blocks["block1_1"] = block1_1
+
+        # stage 2-6
+        for i in range(2, 7):
+            blocks["block%d" % i] = OrderedDict(
+                [
+                    ("Mconv1_stage%d" % i, [150, 128, 7, 1, 3]),
+                    ("Mconv2_stage%d" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv3_stage%d" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv4_stage%d" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv5_stage%d" % i, [128, 128, 7, 1, 3]),
+                    ("Mconv6_stage%d" % i, [128, 128, 1, 1, 0]),
+                    ("Mconv7_stage%d" % i, [128, 22, 1, 1, 0]),
+                ]
+            )
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_0 = blocks["block1_0"]
+        self.model1_1 = blocks["block1_1"]
+        self.model2 = blocks["block2"]
+        self.model3 = blocks["block3"]
+        self.model4 = blocks["block4"]
+        self.model5 = blocks["block5"]
+        self.model6 = blocks["block6"]
+
+    def forward(self, x):
+        out1_0 = self.model1_0(x)
+        out1_1 = self.model1_1(out1_0)
+        concat_stage2 = torch.cat([out1_1, out1_0], 1)
+        out_stage2 = self.model2(concat_stage2)
+        concat_stage3 = torch.cat([out_stage2, out1_0], 1)
+        out_stage3 = self.model3(concat_stage3)
+        concat_stage4 = torch.cat([out_stage3, out1_0], 1)
+        out_stage4 = self.model4(concat_stage4)
+        concat_stage5 = torch.cat([out_stage4, out1_0], 1)
+        out_stage5 = self.model5(concat_stage5)
+        concat_stage6 = torch.cat([out_stage5, out1_0], 1)
+        out_stage6 = self.model6(concat_stage6)
+        return out_stage6
+
+
+class Hand(object):
+    def __init__(self, model_path):
+        self.model = HandPoseModel()
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+        model_dict = transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        scale_search = [0.5, 1.0, 1.5, 2.0]
+        # scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22))
+        # paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = cv2.resize(
+                oriImg,
+                (0, 0),
+                fx=scale,
+                fy=scale,
+                interpolation=cv2.INTER_CUBIC,
+            )
+            imageToTest_padded, pad = padRightDownCorner(
+                imageToTest, stride, padValue
+            )
+            im = (
+                np.transpose(
+                    np.float32(imageToTest_padded[:, :, :, np.newaxis]),
+                    (3, 2, 0, 1),
+                )
+                / 256
+                - 0.5
+            )
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
+            with torch.no_grad():
+                output = self.model(data).cpu().numpy()
+                # output = self.model(data).numpy()q
+
+            # extract outputs, resize, and remove padding
+            heatmap = np.transpose(
+                np.squeeze(output), (1, 2, 0)
+            )  # output 1 is heatmaps
+            heatmap = cv2.resize(
+                heatmap,
+                (0, 0),
+                fx=stride,
+                fy=stride,
+                interpolation=cv2.INTER_CUBIC,
+            )
+            heatmap = heatmap[
+                : imageToTest_padded.shape[0] - pad[2],
+                : imageToTest_padded.shape[1] - pad[3],
+                :,
+            ]
+            heatmap = cv2.resize(
+                heatmap,
+                (oriImg.shape[1], oriImg.shape[0]),
+                interpolation=cv2.INTER_CUBIC,
+            )
+
+            heatmap_avg += heatmap / len(multiplier)
+
+        all_peaks = []
+        for part in range(21):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
+            # 全部小于阈值
+            if np.sum(binary) == 0:
+                all_peaks.append([0, 0])
+                continue
+            label_img, label_numbers = label(
+                binary, return_num=True, connectivity=binary.ndim
+            )
+            max_index = (
+                np.argmax(
+                    [
+                        np.sum(map_ori[label_img == i])
+                        for i in range(1, label_numbers + 1)
+                    ]
+                )
+                + 1
+            )
+            label_img[label_img != max_index] = 0
+            map_ori[label_img == 0] = 0
+
+            y, x = npmax(map_ori)
+            all_peaks.append([x, y])
+        return np.array(all_peaks)

apps/stable_diffusion/src/utils/stencils/openpose/openpose_util.py

@@ -0,0 +1,272 @@
+import math
+import numpy as np
+import matplotlib
+import cv2
+from collections import OrderedDict
+import torch.nn as nn
+
+
+def make_layers(block, no_relu_layers):
+    layers = []
+    for layer_name, v in block.items():
+        if "pool" in layer_name:
+            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])
+            layers.append((layer_name, layer))
+        else:
+            conv2d = nn.Conv2d(
+                in_channels=v[0],
+                out_channels=v[1],
+                kernel_size=v[2],
+                stride=v[3],
+                padding=v[4],
+            )
+            layers.append((layer_name, conv2d))
+            if layer_name not in no_relu_layers:
+                layers.append(("relu_" + layer_name, nn.ReLU(inplace=True)))
+
+    return nn.Sequential(OrderedDict(layers))
+
+
+def padRightDownCorner(img, stride, padValue):
+    h = img.shape[0]
+    w = img.shape[1]
+
+    pad = 4 * [None]
+    pad[0] = 0  # up
+    pad[1] = 0  # left
+    pad[2] = 0 if (h % stride == 0) else stride - (h % stride)  # down
+    pad[3] = 0 if (w % stride == 0) else stride - (w % stride)  # right
+
+    img_padded = img
+    pad_up = np.tile(img_padded[0:1, :, :] * 0 + padValue, (pad[0], 1, 1))
+    img_padded = np.concatenate((pad_up, img_padded), axis=0)
+    pad_left = np.tile(img_padded[:, 0:1, :] * 0 + padValue, (1, pad[1], 1))
+    img_padded = np.concatenate((pad_left, img_padded), axis=1)
+    pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + padValue, (pad[2], 1, 1))
+    img_padded = np.concatenate((img_padded, pad_down), axis=0)
+    pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + padValue, (1, pad[3], 1))
+    img_padded = np.concatenate((img_padded, pad_right), axis=1)
+
+    return img_padded, pad
+
+
+# transfer caffe model to pytorch which will match the layer name
+def transfer(model, model_weights):
+    transfered_model_weights = {}
+    for weights_name in model.state_dict().keys():
+        transfered_model_weights[weights_name] = model_weights[
+            ".".join(weights_name.split(".")[1:])
+        ]
+    return transfered_model_weights
+
+
+# draw the body keypoint and lims
+def draw_bodypose(canvas, candidate, subset):
+    stickwidth = 4
+    limbSeq = [
+        [2, 3],
+        [2, 6],
+        [3, 4],
+        [4, 5],
+        [6, 7],
+        [7, 8],
+        [2, 9],
+        [9, 10],
+        [10, 11],
+        [2, 12],
+        [12, 13],
+        [13, 14],
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],
+        [3, 17],
+        [6, 18],
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+    ]
+    for i in range(18):
+        for n in range(len(subset)):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
+    for i in range(17):
+        for n in range(len(subset)):
+            index = subset[n][np.array(limbSeq[i]) - 1]
+            if -1 in index:
+                continue
+            cur_canvas = canvas.copy()
+            Y = candidate[index.astype(int), 0]
+            X = candidate[index.astype(int), 1]
+            mX = np.mean(X)
+            mY = np.mean(Y)
+            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+            polygon = cv2.ellipse2Poly(
+                (int(mY), int(mX)),
+                (int(length / 2), stickwidth),
+                int(angle),
+                0,
+                360,
+                1,
+            )
+            cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
+            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
+    return canvas
+
+
+# image drawed by opencv is not good.
+def draw_handpose(canvas, all_hand_peaks, show_number=False):
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for peaks in all_hand_peaks:
+        for ie, e in enumerate(edges):
+            if np.sum(np.all(peaks[e], axis=1) == 0) == 0:
+                x1, y1 = peaks[e[0]]
+                x2, y2 = peaks[e[1]]
+                cv2.line(
+                    canvas,
+                    (x1, y1),
+                    (x2, y2),
+                    matplotlib.colors.hsv_to_rgb(
+                        [ie / float(len(edges)), 1.0, 1.0]
+                    )
+                    * 255,
+                    thickness=2,
+                )
+
+        for i, keyponit in enumerate(peaks):
+            x, y = keyponit
+            cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+            if show_number:
+                cv2.putText(
+                    canvas,
+                    str(i),
+                    (x, y),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.3,
+                    (0, 0, 0),
+                    lineType=cv2.LINE_AA,
+                )
+    return canvas
+
+
+# detect hand according to body pose keypoints
+# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
+def handDetect(candidate, subset, oriImg):
+    # right hand: wrist 4, elbow 3, shoulder 2
+    # left hand: wrist 7, elbow 6, shoulder 5
+    ratioWristElbow = 0.33
+    detect_result = []
+    image_height, image_width = oriImg.shape[0:2]
+    for person in subset.astype(int):
+        # if any of three not detected
+        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
+        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
+        if not (has_left or has_right):
+            continue
+        hands = []
+        # left hand
+        if has_left:
+            left_shoulder_index, left_elbow_index, left_wrist_index = person[
+                [5, 6, 7]
+            ]
+            x1, y1 = candidate[left_shoulder_index][:2]
+            x2, y2 = candidate[left_elbow_index][:2]
+            x3, y3 = candidate[left_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, True])
+        # right hand
+        if has_right:
+            (
+                right_shoulder_index,
+                right_elbow_index,
+                right_wrist_index,
+            ) = person[[2, 3, 4]]
+            x1, y1 = candidate[right_shoulder_index][:2]
+            x2, y2 = candidate[right_elbow_index][:2]
+            x3, y3 = candidate[right_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, False])
+
+        for x1, y1, x2, y2, x3, y3, is_left in hands:
+            x = x3 + ratioWristElbow * (x3 - x2)
+            y = y3 + ratioWristElbow * (y3 - y2)
+            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
+            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
+            # x-y refers to the center --> offset to topLeft point
+            x -= width / 2
+            y -= width / 2  # width = height
+            # overflow the image
+            if x < 0:
+                x = 0
+            if y < 0:
+                y = 0
+            width1 = width
+            width2 = width
+            if x + width > image_width:
+                width1 = image_width - x
+            if y + width > image_height:
+                width2 = image_height - y
+            width = min(width1, width2)
+            # the max hand box value is 20 pixels
+            if width >= 20:
+                detect_result.append([int(x), int(y), int(width), is_left])
+
+    """
+    return value: [[x, y, w, True if left hand else False]].
+    width=height since the network require squared input.
+    x, y is the coordinate of top left 
+    """
+    return detect_result
+
+
+# get max index of 2d array
+def npmax(array):
+    arrayindex = array.argmax(1)
+    arrayvalue = array.max(1)
+    i = arrayvalue.argmax()
+    j = arrayindex[i]
+    return (i,)

apps/stable_diffusion/src/utils/stencils/stencil_utils.py

@@ -1,8 +1,10 @@
-import cv2
 import numpy as np
 from PIL import Image
 import torch
-from apps.stable_diffusion.src.utils.stencils.canny import CannyDetector
+from apps.stable_diffusion.src.utils.stencils import (
+    CannyDetector,
+    OpenposeDetector,
+)
 
 stencil = {}
 
@@ -26,23 +28,6 @@ def HWC3(x):
         return y
 
 
-def resize_image(input_image, resolution):
-    H, W, C = input_image.shape
-    H = float(H)
-    W = float(W)
-    k = float(resolution) / min(H, W)
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    img = cv2.resize(
-        input_image,
-        (W, H),
-        interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA,
-    )
-    return img
-
-
 def controlnet_hint_shaping(
     controlnet_hint, height, width, dtype, num_images_per_prompt=1
 ):
@@ -125,7 +110,13 @@ def controlnet_hint_conversion(
     match use_stencil:
         case "canny":
             print("Detecting edge with canny")
-            controlnet_hint = hint_canny(image, width)
+            controlnet_hint = hint_canny(image)
+        case "openpose":
+            print("Detecting human pose")
+            controlnet_hint = hint_openpose(image)
+        case "scribble":
+            print("Working with scribble")
+            controlnet_hint = hint_scribble(image)
         case _:
             return None
     controlnet_hint = controlnet_hint_shaping(
@@ -134,22 +125,62 @@ def controlnet_hint_conversion(
     return controlnet_hint
 
 
+stencil_to_model_id_map = {
+    "canny": "lllyasviel/sd-controlnet-canny",
+    "depth": "lllyasviel/sd-controlnet-depth",
+    "hed": "lllyasviel/sd-controlnet-hed",
+    "mlsd": "lllyasviel/sd-controlnet-mlsd",
+    "normal": "lllyasviel/sd-controlnet-normal",
+    "openpose": "lllyasviel/sd-controlnet-openpose",
+    "scribble": "lllyasviel/sd-controlnet-scribble",
+    "seg": "lllyasviel/sd-controlnet-seg",
+}
+
+
+def get_stencil_model_id(use_stencil):
+    if use_stencil in stencil_to_model_id_map:
+        return stencil_to_model_id_map[use_stencil]
+    return None
+
+
 # Stencil 1. Canny
 def hint_canny(
     image: Image.Image,
-    width=512,
-    height=512,
     low_threshold=100,
     high_threshold=200,
 ):
     with torch.no_grad():
         input_image = np.array(image)
-        image_resolution = width
-
-        img = resize_image(HWC3(input_image), image_resolution)
 
         if not "canny" in stencil:
             stencil["canny"] = CannyDetector()
-        detected_map = stencil["canny"](img, low_threshold, high_threshold)
+        detected_map = stencil["canny"](
+            input_image, low_threshold, high_threshold
+        )
         detected_map = HWC3(detected_map)
         return detected_map
+
+
+# Stencil 2. OpenPose.
+def hint_openpose(
+    image: Image.Image,
+):
+    with torch.no_grad():
+        input_image = np.array(image)
+
+        if not "openpose" in stencil:
+            stencil["openpose"] = OpenposeDetector()
+
+        detected_map, _ = stencil["openpose"](input_image)
+        detected_map = HWC3(detected_map)
+        return detected_map
+
+
+# Stencil 3. Scribble.
+def hint_scribble(image: Image.Image):
+    with torch.no_grad():
+        input_image = np.array(image)
+
+        detected_map = np.zeros_like(input_image, dtype=np.uint8)
+        detected_map[np.min(input_image, axis=2) < 127] = 255
+        return detected_map

apps/stable_diffusion/src/utils/utils.py

@@ -8,6 +8,7 @@ from csv import DictWriter
 from pathlib import Path
 import numpy as np
 from random import randint
+import tempfile
 from shark.shark_inference import SharkInference
 from shark.shark_importer import import_with_fx
 from shark.iree_utils.vulkan_utils import (
@@ -25,11 +26,7 @@ from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
 
 
 def get_extended_name(model_name):
-    device = (
-        args.device
-        if "://" not in args.device
-        else "-".join(args.device.split("://"))
-    )
+    device = args.device.split("://", 1)[0]
     extended_name = "{}_{}".format(model_name, device)
     return extended_name
 
@@ -94,6 +91,9 @@ def compile_through_fx(
     is_f16=False,
     f16_input_mask=None,
     use_tuned=False,
+    save_dir=tempfile.gettempdir(),
+    debug=False,
+    generate_vmfb=True,
     extra_args=[],
 ):
     from shark.parser import shark_args
@@ -101,10 +101,18 @@ def compile_through_fx(
     if "cuda" in args.device:
         shark_args.enable_tf32 = True
 
-    mlir_module, func_name = import_with_fx(
-        model, inputs, is_f16, f16_input_mask
+    (
+        mlir_module,
+        func_name,
+    ) = import_with_fx(
+        model=model,
+        inputs=inputs,
+        is_f16=is_f16,
+        f16_input_mask=f16_input_mask,
+        debug=debug,
+        model_name=model_name,
+        save_dir=save_dir,
     )
-
     if use_tuned:
         if "vae" in model_name.split("_")[0]:
             args.annotation_model = "vae"
@@ -116,11 +124,19 @@ def compile_through_fx(
         mlir_dialect="linalg",
     )
 
+    if generate_vmfb:
+        shark_module = SharkInference(
+            mlir_module,
+            device=args.device,
+            mlir_dialect="linalg",
+        )
+        del mlir_module
+        gc.collect()
+        return _compile_module(shark_module, model_name, extra_args)
+
     del mlir_module
     gc.collect()
 
-    return _compile_module(shark_module, model_name, extra_args)
-
 
 def set_iree_runtime_flags():
     vulkan_runtime_flags = [
@@ -563,7 +579,7 @@ def save_output_img(output_img, img_seed, extra_info={}):
 
     img_model = args.hf_model_id
     if args.ckpt_loc:
-        img_model = os.path.basename(args.ckpt_loc)
+        img_model = Path(os.path.basename(args.ckpt_loc)).stem
 
     if args.output_img_format == "jpg":
         out_img_path = Path(generated_imgs_path, f"{out_img_name}.jpg")
@@ -603,7 +619,7 @@ def save_output_img(output_img, img_seed, extra_info={}):
 
     new_entry.update(extra_info)
 
-    with open(csv_path, "a") as csv_obj:
+    with open(csv_path, "a", encoding="utf-8") as csv_obj:
         dictwriter_obj = DictWriter(csv_obj, fieldnames=list(new_entry.keys()))
         dictwriter_obj.writerow(new_entry)
         csv_obj.close()

apps/stable_diffusion/web/index.py

@@ -9,9 +9,12 @@ from apps.stable_diffusion.src import args, clear_all
 from apps.stable_diffusion.web.utils.gradio_configs import (
     clear_gradio_tmp_imgs_folder,
 )
+from apps.stable_diffusion.web.ui.utils import get_custom_model_path
 
-# clear all gradio tmp images from the last session
+# Clear all gradio tmp images from the last session
 clear_gradio_tmp_imgs_folder()
+# Create the custom model folder if it doesn't already exist
+get_custom_model_path().mkdir(parents=True, exist_ok=True)
 
 if args.clear_all:
     clear_all()

apps/stable_diffusion/web/ui/css/sd_dark_theme.css

@@ -1,147 +1,93 @@
-
-/* Overwrite the Gradio default theme with their .dark theme declarations */
+/*
+Apply Gradio dark theme to the default Gradio theme.
+Procedure to upgrade the dark theme:
+- Using your browser, visit http://localhost:8080/?__theme=dark
+- Open your browser inspector, search for the .dark css class
+- Copy .dark class declarations, apply them here into :root
+*/
 
 :root {
-    --color-focus-primary: var(--color-grey-700);
-    --color-focus-secondary: var(--color-grey-600);
-    --color-focus-ring: rgb(55 65 81);
-    --color-background-primary: var(--color-grey-950);
-    --color-background-secondary: var(--color-grey-900);
-    --color-background-tertiary: var(--color-grey-800);
-    --color-text-body: var(--color-grey-100);
-    --color-text-label: var(--color-grey-200);
-    --color-text-placeholder: var(--color-grey);
-    --color-text-subdued: var(--color-grey-400);
-    --color-text-link-base: var(--color-blue-500);
-    --color-text-link-hover: var(--color-blue-400);
-    --color-text-link-visited: var(--color-blue-600);
-    --color-text-link-active: var(--color-blue-500);
-    --color-text-code-background: var(--color-grey-800);
-    --color-text-code-border: color.border-primary;
-    --color-border-primary: var(--color-grey-700);
-    --color-border-secondary: var(--color-grey-600);
-    --color-border-highlight: var(--color-accent-base);
-    --color-accent-base: var(--color-orange-500);
-    --color-accent-light: var(--color-orange-300);
-    --color-accent-dark: var(--color-orange-700);
-    --color-functional-error-base: var(--color-red-400);
-    --color-functional-error-subdued: var(--color-red-300);
-    --color-functional-error-background: var(--color-background-primary);
-    --color-functional-info-base: var(--color-yellow);
-    --color-functional-info-subdued: var(--color-yellow-300);
-    --color-functional-success-base: var(--color-green);
-    --color-functional-success-subdued: var(--color-green-300);
-    --shadow-spread: 2px;
-    --api-background: linear-gradient(to bottom, rgba(255, 216, 180, .05), transparent);
-    --api-pill-background: var(--color-orange-400);
-    --api-pill-border: var(--color-orange-600);
-    --api-pill-text: var(--color-orange-900);
+    --color-accent-soft: var(--neutral-700);
+    --color-background-primary: var(--neutral-950);
+    --color-background-secondary: var(--neutral-900);
+    --color-border-accent: var(--neutral-600);
+    --color-border-primary: var(--neutral-700);
+    --text-color-code-background: var(--neutral-800);
+    --text-color-link-active: var(--secondary-500);
+    --text-color-link: var(--secondary-500);
+    --text-color-link-hover: var(--secondary-400);
+    --text-color-link-visited: var(--secondary-600);
+    --text-color-subdued: var(--neutral-400);
+    --body-background-color: var(--color-background-primary);
+    --body-text-color: var(--neutral-100);
+    --shadow-spread: 1px;
+    --block-background: var(--neutral-800);
     --block-border-color: var(--color-border-primary);
-    --block-background: var(--color-background-tertiary);
-    --uploadable-border-color-hover: var(--color-border-primary);
-    --uploadable-border-color-loaded: var(--color-functional-success);
-    --uploadable-text-color: var(--color-text-subdued);
-    --block_label-border-color: var(--color-border-primary);
-    --block_label-icon-color: var(--color-text-label);
-    --block_label-shadow: var(--shadow-drop);
-    --block_label-background: var(--color-background-secondary);
-    --icon_button-icon-color-base: var(--color-text-label);
-    --icon_button-icon-color-hover: var(--color-text-label);
-    --icon_button-background-base: var(--color-background-primary);
-    --icon_button-background-hover: var(--color-background-primary);
-    --icon_button-border-color-base: var(--color-background-primary);
-    --icon_button-border-color-hover: var(--color-border-secondary);
-    --input-text-color: var(--color-text-body);
-    --input-border-color-base: var(--color-border-primary);
-    --input-border-color-hover: var(--color-border-primary);
-    --input-border-color-focus: var(--color-border-primary);
-    --input-background-base: var(--color-background-tertiary);
-    --input-background-hover: var(--color-background-tertiary);
-    --input-background-focus: var(--color-background-tertiary);
-    --input-shadow: var(--shadow-inset);
-    --checkbox-border-color-base: var(--color-border-primary);
-    --checkbox-border-color-hover: var(--color-focus-primary);
-    --checkbox-border-color-focus: var(--color-blue-500);
-    --checkbox-background-base: var(--color-background-primary);
-    --checkbox-background-hover: var(--color-background-primary);
-    --checkbox-background-focus: var(--color-background-primary);
-    --checkbox-background-selected: var(--color-blue-600);
-    --checkbox-label-border-color-base: var(--color-border-primary);
+    --block-border-width: 1px;
+    --block-info-color: var(--text-color-subdued);
+    --block-label-background: var(--color-background-secondary);
+    --block-label-border-color: var(--color-border-primary);
+    --block-label-border-width: 1px;
+    --block-label-color: var(--neutral-200);
+    --block-shadow: none;
+    --block-title-background: none;
+    --block-title-border-color: none;
+    --block-title-border-width: 0px;
+    --block-title-color: var(--neutral-200);
+    --panel-background: var(--color-background-secondary);
+    --panel-border-color: var(--color-border-primary);
+    --checkbox-background: var(--neutral-800);
+    --checkbox-background-focus: var(--checkbox-background);
+    --checkbox-background-hover: var(--checkbox-background);
+    --checkbox-background-selected: var(--secondary-600);
+    --checkbox-border-color: var(--neutral-700);
+    --checkbox-border-color-focus: var(--secondary-500);
+    --checkbox-border-color-hover: var(--neutral-600);
+    --checkbox-border-color-selected: var(--secondary-600);
+    --checkbox-label-background: linear-gradient(to top, var(--neutral-900), var(--neutral-800));
+    --checkbox-label-background-hover: linear-gradient(to top, var(--neutral-900), var(--neutral-800));
+    --checkbox-label-background-selected: var(--checkbox-label-background);
+    --checkbox-label-border-color: var(--color-border-primary);
     --checkbox-label-border-color-hover: var(--color-border-primary);
-    --checkbox-label-border-color-focus: var(--color-border-secondary);
-    --checkbox-label-background-base: linear-gradient(to top, var(--color-grey-900), var(--color-grey-800));
-    --checkbox-label-background-hover: linear-gradient(to top, var(--color-grey-900), var(--color-grey-800));
-    --checkbox-label-background-focus: linear-gradient(to top, var(--color-grey-900), var(--color-grey-800));
-    --form-seperator-color: var(--color-border-primary);
-    --button-primary-border-color-base: var(--color-orange-600);
-    --button-primary-border-color-hover: var(--color-orange-600);
-    --button-primary-border-color-focus: var(--color-orange-600);
-    --button-primary-text-color-base: white;
-    --button-primary-text-color-hover: white;
-    --button-primary-text-color-focus: white;
-    --button-primary-background-base: linear-gradient(to bottom right, var(--color-orange-700), var(--color-orange-700));
-    --button-primary-background-hover: linear-gradient(to bottom right, var(--color-orange-700), var(--color-orange-500));
-    --button-primary-background-focus: linear-gradient(to bottom right, var(--color-orange-700), var(--color-orange-500));
-    --button-secondary-border-color-base: var(--color-grey-600);
-    --button-secondary-border-color-hover: var(--color-grey-600);
-    --button-secondary-border-color-focus: var(--color-grey-600);
-    --button-secondary-text-color-base: white;
-    --button-secondary-text-color-hover: white;
-    --button-secondary-text-color-focus: white;
-    --button-secondary-background-base: linear-gradient(to bottom right, var(--color-grey-600), var(--color-grey-700));
-    --button-secondary-background-hover: linear-gradient(to bottom right, var(--color-grey-600), var(--color-grey-600));
-    --button-secondary-background-focus: linear-gradient(to bottom right, var(--color-grey-600), var(--color-grey-600));
-    --button-cancel-border-color-base: var(--color-red-600);
-    --button-cancel-border-color-hover: var(--color-red-600);
-    --button-cancel-border-color-focus: var(--color-red-600);
-    --button-cancel-text-color-base: white;
-    --button-cancel-text-color-hover: white;
-    --button-cancel-text-color-focus: white;
-    --button-cancel-background-base: linear-gradient(to bottom right, var(--color-red-700), var(--color-red-700));
-    --button-cancel-background-focus: linear-gradient(to bottom right, var(--color-red-700), var(--color-red-500));
-    --button-cancel-background-hover: linear-gradient(to bottom right, var(--color-red-700), var(--color-red-500));
-    --button-plain-border-color-base: var(--color-grey-600);
-    --button-plain-border-color-hover: var(--color-grey-500);
-    --button-plain-border-color-focus: var(--color-grey-500);
-    --button-plain-text-color-base: var(--color-text-body);
-    --button-plain-text-color-hover: var(--color-text-body);
-    --button-plain-text-color-focus: var(--color-text-body);
-    --button-plain-background-base: var(--color-grey-700);
-    --button-plain-background-hover: var(--color-grey-700);
-    --button-plain-background-focus: var(--color-grey-700);
-    --gallery-label-background-base: var(--color-grey-50);
-    --gallery-label-background-hover: var(--color-grey-50);
-    --gallery-label-border-color-base: var(--color-border-primary);
-    --gallery-label-border-color-hover: var(--color-border-primary);
-    --gallery-thumb-background-base: var(--color-grey-900);
-    --gallery-thumb-background-hover: var(--color-grey-900);
-    --gallery-thumb-border-color-base: var(--color-border-primary);
-    --gallery-thumb-border-color-hover: var(--color-accent-base);
-    --gallery-thumb-border-color-focus: var(--color-blue-500);
-    --gallery-thumb-border-color-selected: var(--color-accent-base);
-    --chatbot-border-border-color-base: transparent;
-    --chatbot-border-border-color-latest: transparent;
-    --chatbot-user-background-base: ;
-    --chatbot-user-background-latest: ;
-    --chatbot-user-text-color-base: white;
-    --chatbot-user-text-color-latest: white;
-    --chatbot-bot-background-base: ;
-    --chatbot-bot-background-latest: ;
-    --chatbot-bot-text-color-base: white;
-    --chatbot-bot-text-color-latest: white;
-    --label-gradient-from: var(--color-orange-400);
-    --label-gradient-to: var(--color-orange-600);
-    --table-odd-background: var(--color-grey-900);
-    --table-even-background: var(--color-grey-950);
-    --table-background-edit: transparent;
-    --dataset-gallery-background-base: var(--color-background-primary);
-    --dataset-gallery-background-hover: var(--color-grey-800);
-    --dataset-dataframe-border-base: var(--color-border-primary);
-    --dataset-dataframe-border-hover: var(--color-border-secondary);
-    --dataset-table-background-base: transparent;
-    --dataset-table-background-hover: var(--color-grey-700);
-    --dataset-table-border-base: var(--color-grey-800);
-    --dataset-table-border-hover: var(--color-grey-800);
+    --checkbox-text-color: var(--body-text-color);
+    --checkbox-text-color-selected: var(--checkbox-text-color);
+    --error-background: var(--color-background-primary);
+    --error-border-color: var(--color-border-primary);
+    --error-border-width: var(--error-border-width);
+    --error-color: #ef4444;
+    --input-background: var(--neutral-800);
+    --input-background-focus: var(--secondary-600);
+    --input-background-hover: var(--input-background);
+    --input-border-color: var(--color-border-primary);
+    --input-border-color-focus: var(--neutral-700);
+    --input-border-color-hover: var(--color-border-primary);
+    --input-placeholder-color: var(--neutral-500);
+    --input-shadow: var(--input-shadow);
+    --input-shadow-focus: 0 0 0 var(--shadow-spread) var(--neutral-700), var(--shadow-inset);
+    --loader-color: var(--color-accent);
+    --stat-color-background: linear-gradient(to right, var(--primary-400), var(--primary-600));
+    --table-border-color: var(--neutral-700);
+    --table-even-background: var(--neutral-950);
+    --table-odd-background: var(--neutral-900);
+    --table-row-focus: var(--color-accent-soft);
+    --button-cancel-background: linear-gradient(to bottom right, #dc2626, #b91c1c);
+    --button-cancel-background-hover: linear-gradient(to bottom right, #dc2626, #dc2626);
+    --button-cancel-border-color: #dc2626;
+    --button-cancel-border-color-hover: var(--button-cancel-border-color);
+    --button-cancel-text-color: white;
+    --button-cancel-text-color-hover: var(--button-cancel-text-color);
+    --button-primary-background: linear-gradient(to bottom right, var(--primary-600), var(--primary-700));
+    --button-primary-background-hover: linear-gradient(to bottom right, var(--primary-600), var(--primary-600));
+    --button-primary-border-color: var(--primary-600);
+    --button-primary-border-color-hover: var(--button-primary-border-color);
+    --button-primary-text-color: white;
+    --button-primary-text-color-hover: var(--button-primary-text-color);
+    --button-secondary-background: linear-gradient(to bottom right, var(--neutral-600), var(--neutral-700));
+    --button-secondary-background-hover: linear-gradient(to bottom right, var(--neutral-600), var(--neutral-600));
+    --button-secondary-border-color: var(--neutral-600);
+    --button-secondary-border-color-hover: var(--button-secondary-border-color);
+    --button-secondary-text-color: white;
+    --button-secondary-text-color-hover: var(--button-secondary-text-color);
 }
 
 /* SHARK theme */
@@ -219,3 +165,28 @@ footer {
     pointer-events: none;
 }
 
+/* Import Png info box */
+#txt2img_prompt_image .fixed-height {
+    height: var(--size-32);
+}
+
+/* Hide "remove buttons" from ui dropdowns */
+#custom_model .token-remove.remove-all,
+#scheduler .token-remove.remove-all,
+#device .token-remove.remove-all,
+#stencil_model .token-remove.remove-all {
+    display: none;
+}
+
+/* Hide selected items from ui dropdowns */
+#custom_model .options .item .inner-item,
+#scheduler .options .item .inner-item,
+#device .options .item .inner-item,
+#stencil_model .options .item .inner-item {
+    display:none;
+}
+
+/* Hide the download icon from the nod logo */
+#top_logo .download {
+    display: none;
+}

apps/stable_diffusion/web/ui/img2img_ui.py

@@ -1,7 +1,5 @@
-import os
-import sys
-import glob
 from pathlib import Path
+import os
 import gradio as gr
 from PIL import Image
 from apps.stable_diffusion.scripts import img2img_inf
@@ -9,6 +7,10 @@ from apps.stable_diffusion.src import args
 from apps.stable_diffusion.web.ui.utils import (
     available_devices,
     nodlogo_loc,
+    get_custom_model_path,
+    get_custom_model_files,
+    scheduler_list,
+    predefined_models,
 )
 
 
@@ -27,34 +29,18 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
         with gr.Row():
             with gr.Column(scale=1, min_width=600):
                 with gr.Row():
-                    ckpt_path = (
-                        Path(args.ckpt_dir)
-                        if args.ckpt_dir
-                        else Path(Path.cwd(), "models")
-                    )
-                    ckpt_path.mkdir(parents=True, exist_ok=True)
-                    types = (
-                        "*.ckpt",
-                        "*.safetensors",
-                    )  # the tuple of file types
-                    ckpt_files = ["None"]
-                    for extn in types:
-                        files = glob.glob(os.path.join(ckpt_path, extn))
-                        ckpt_files.extend(files)
                     custom_model = gr.Dropdown(
-                        label=f"Models (Custom Model path: {ckpt_path})",
-                        value=args.ckpt_loc if args.ckpt_loc else "None",
-                        choices=ckpt_files
-                        + [
-                            "Linaqruf/anything-v3.0",
-                            "prompthero/openjourney",
-                            "wavymulder/Analog-Diffusion",
-                            "stabilityai/stable-diffusion-2-1",
-                            "stabilityai/stable-diffusion-2-1-base",
-                            "CompVis/stable-diffusion-v1-4",
-                        ],
+                        label=f"Models (Custom Model path: {get_custom_model_path()})",
+                        elem_id="custom_model",
+                        value=os.path.basename(args.ckpt_loc)
+                        if args.ckpt_loc
+                        else "None",
+                        choices=["None"]
+                        + get_custom_model_files()
+                        + predefined_models,
                     )
                     hf_model_id = gr.Textbox(
+                        elem_id="hf_model_id",
                         placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3",
                         value="",
                         label="HuggingFace Model ID",
@@ -82,21 +68,18 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                 with gr.Accordion(label="Stencil Options", open=False):
                     with gr.Row():
                         use_stencil = gr.Dropdown(
+                            elem_id="stencil_model",
                             label="Stencil model",
                             value="None",
-                            choices=["None", "canny"],
+                            choices=["None", "canny", "openpose", "scribble"],
                         )
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():
                         scheduler = gr.Dropdown(
+                            elem_id="scheduler",
                             label="Scheduler",
                             value="PNDM",
-                            choices=[
-                                "DDIM",
-                                "PNDM",
-                                "DPMSolverMultistep",
-                                "EulerAncestralDiscrete",
-                            ],
+                            choices=scheduler_list,
                         )
                         with gr.Group():
                             save_metadata_to_png = gr.Checkbox(
@@ -143,7 +126,7 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                             1,
                             value=args.strength,
                             step=0.01,
-                            label="Strength",
+                            label="Denoising Strength",
                         )
                     with gr.Row():
                         guidance_scale = gr.Slider(
@@ -175,19 +158,24 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                         value=args.seed, precision=0, label="Seed"
                     )
                     device = gr.Dropdown(
+                        elem_id="device",
                         label="Device",
                         value=available_devices[0],
                         choices=available_devices,
                     )
                 with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        None,
-                        inputs=[],
-                        outputs=[seed],
-                        _js="() => Math.floor(Math.random() * 4294967295)",
-                    )
-                    stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=2):
+                        random_seed = gr.Button("Randomize Seed")
+                        random_seed.click(
+                            None,
+                            inputs=[],
+                            outputs=[seed],
+                            _js="() => -1",
+                        )
+                    with gr.Column(scale=6):
+                        stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=1, min_width=150):
+                        clear_queue = gr.Button("Clear Queue")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -242,6 +230,9 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
             show_progress=args.progress_bar,
         )
 
-        prompt.submit(**kwargs)
-        negative_prompt.submit(**kwargs)
-        stable_diffusion.click(**kwargs)
+        prompt_submit = prompt.submit(**kwargs)
+        neg_prompt_submit = negative_prompt.submit(**kwargs)
+        generate_click = stable_diffusion.click(**kwargs)
+        clear_queue.click(
+            fn=None, cancels=[prompt_submit, neg_prompt_submit, generate_click]
+        )

apps/stable_diffusion/web/ui/inpaint_ui.py

@@ -1,7 +1,5 @@
-import os
-import sys
-import glob
 from pathlib import Path
+import os
 import gradio as gr
 from PIL import Image
 from apps.stable_diffusion.scripts import inpaint_inf
@@ -9,6 +7,10 @@ from apps.stable_diffusion.src import args
 from apps.stable_diffusion.web.ui.utils import (
     available_devices,
     nodlogo_loc,
+    get_custom_model_path,
+    get_custom_model_files,
+    scheduler_list,
+    predefined_paint_models,
 )
 
 
@@ -27,30 +29,18 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
         with gr.Row():
             with gr.Column(scale=1, min_width=600):
                 with gr.Row():
-                    ckpt_path = (
-                        Path(args.ckpt_dir)
-                        if args.ckpt_dir
-                        else Path(Path.cwd(), "models")
-                    )
-                    ckpt_path.mkdir(parents=True, exist_ok=True)
-                    types = (
-                        "*.ckpt",
-                        "*.safetensors",
-                    )  # the tuple of file types
-                    ckpt_files = ["None"]
-                    for extn in types:
-                        files = glob.glob(os.path.join(ckpt_path, extn))
-                        ckpt_files.extend(files)
                     custom_model = gr.Dropdown(
-                        label=f"Models (Custom Model path: {ckpt_path})",
-                        value=args.ckpt_loc if args.ckpt_loc else "None",
-                        choices=ckpt_files
-                        + [
-                            "runwayml/stable-diffusion-inpainting",
-                            "stabilityai/stable-diffusion-2-inpainting",
-                        ],
+                        label=f"Models (Custom Model path: {get_custom_model_path()})",
+                        elem_id="custom_model",
+                        value=os.path.basename(args.ckpt_loc)
+                        if args.ckpt_loc
+                        else "None",
+                        choices=["None"]
+                        + get_custom_model_files()
+                        + predefined_paint_models,
                     )
                     hf_model_id = gr.Textbox(
+                        elem_id="hf_model_id",
                         placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: ghunkins/stable-diffusion-liberty-inpainting",
                         value="",
                         label="HuggingFace Model ID",
@@ -81,14 +71,10 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():
                         scheduler = gr.Dropdown(
+                            elem_id="scheduler",
                             label="Scheduler",
                             value="PNDM",
-                            choices=[
-                                "DDIM",
-                                "PNDM",
-                                "DPMSolverMultistep",
-                                "EulerAncestralDiscrete",
-                            ],
+                            choices=scheduler_list,
                         )
                         with gr.Group():
                             save_metadata_to_png = gr.Checkbox(
@@ -126,6 +112,20 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                             ],
                             visible=False,
                         )
+                    with gr.Row():
+                        inpaint_full_res = gr.Radio(
+                            choices=["Whole picture", "Only masked"],
+                            type="index",
+                            value="Whole picture",
+                            label="Inpaint area",
+                        )
+                        inpaint_full_res_padding = gr.Slider(
+                            minimum=0,
+                            maximum=256,
+                            step=4,
+                            value=32,
+                            label="Only masked padding, pixels",
+                        )
                     with gr.Row():
                         steps = gr.Slider(
                             1, 100, value=args.steps, step=1, label="Steps"
@@ -160,19 +160,24 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                         value=args.seed, precision=0, label="Seed"
                     )
                     device = gr.Dropdown(
+                        elem_id="device",
                         label="Device",
                         value=available_devices[0],
                         choices=available_devices,
                     )
                 with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        None,
-                        inputs=[],
-                        outputs=[seed],
-                        _js="() => Math.floor(Math.random() * 4294967295)",
-                    )
-                    stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=2):
+                        random_seed = gr.Button("Randomize Seed")
+                        random_seed.click(
+                            None,
+                            inputs=[],
+                            outputs=[seed],
+                            _js="() => -1",
+                        )
+                    with gr.Column(scale=6):
+                        stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=1, min_width=150):
+                        clear_queue = gr.Button("Clear Queue")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -207,6 +212,8 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                 inpaint_init_image,
                 height,
                 width,
+                inpaint_full_res,
+                inpaint_full_res_padding,
                 steps,
                 guidance_scale,
                 seed,
@@ -225,6 +232,9 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
             show_progress=args.progress_bar,
         )
 
-        prompt.submit(**kwargs)
-        negative_prompt.submit(**kwargs)
-        stable_diffusion.click(**kwargs)
+        prompt_submit = prompt.submit(**kwargs)
+        neg_prompt_submit = negative_prompt.submit(**kwargs)
+        generate_click = stable_diffusion.click(**kwargs)
+        clear_queue.click(
+            fn=None, cancels=[prompt_submit, neg_prompt_submit, generate_click]
+        )

apps/stable_diffusion/web/ui/outpaint_ui.py

@@ -1,7 +1,5 @@
-import os
-import sys
-import glob
 from pathlib import Path
+import os
 import gradio as gr
 from PIL import Image
 from apps.stable_diffusion.scripts import outpaint_inf
@@ -9,6 +7,10 @@ from apps.stable_diffusion.src import args
 from apps.stable_diffusion.web.ui.utils import (
     available_devices,
     nodlogo_loc,
+    get_custom_model_path,
+    get_custom_model_files,
+    scheduler_list,
+    predefined_paint_models,
 )
 
 
@@ -27,30 +29,18 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
         with gr.Row():
             with gr.Column(scale=1, min_width=600):
                 with gr.Row():
-                    ckpt_path = (
-                        Path(args.ckpt_dir)
-                        if args.ckpt_dir
-                        else Path(Path.cwd(), "models")
-                    )
-                    ckpt_path.mkdir(parents=True, exist_ok=True)
-                    types = (
-                        "*.ckpt",
-                        "*.safetensors",
-                    )  # the tuple of file types
-                    ckpt_files = ["None"]
-                    for extn in types:
-                        files = glob.glob(os.path.join(ckpt_path, extn))
-                        ckpt_files.extend(files)
                     custom_model = gr.Dropdown(
-                        label=f"Models (Custom Model path: {ckpt_path})",
-                        value=args.ckpt_loc if args.ckpt_loc else "None",
-                        choices=ckpt_files
-                        + [
-                            "runwayml/stable-diffusion-inpainting",
-                            "stabilityai/stable-diffusion-2-inpainting",
-                        ],
+                        label=f"Models (Custom Model path: {get_custom_model_path()})",
+                        elem_id="custom_model",
+                        value=os.path.basename(args.ckpt_loc)
+                        if args.ckpt_loc
+                        else "None",
+                        choices=["None"]
+                        + get_custom_model_files()
+                        + predefined_paint_models,
                     )
                     hf_model_id = gr.Textbox(
+                        elem_id="hf_model_id",
                         placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: ghunkins/stable-diffusion-liberty-inpainting",
                         value="",
                         label="HuggingFace Model ID",
@@ -78,14 +68,10 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():
                         scheduler = gr.Dropdown(
+                            elem_id="scheduler",
                             label="Scheduler",
                             value="PNDM",
-                            choices=[
-                                "DDIM",
-                                "PNDM",
-                                "DPMSolverMultistep",
-                                "EulerAncestralDiscrete",
-                            ],
+                            choices=scheduler_list,
                         )
                         with gr.Group():
                             save_metadata_to_png = gr.Checkbox(
@@ -193,19 +179,24 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
                         value=args.seed, precision=0, label="Seed"
                     )
                     device = gr.Dropdown(
+                        elem_id="device",
                         label="Device",
                         value=available_devices[0],
                         choices=available_devices,
                     )
                 with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        None,
-                        inputs=[],
-                        outputs=[seed],
-                        _js="() => Math.floor(Math.random() * 4294967295)",
-                    )
-                    stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=2):
+                        random_seed = gr.Button("Randomize Seed")
+                        random_seed.click(
+                            None,
+                            inputs=[],
+                            outputs=[seed],
+                            _js="() => -1",
+                        )
+                    with gr.Column(scale=6):
+                        stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=1, min_width=150):
+                        clear_queue = gr.Button("Clear Queue")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -261,6 +252,9 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
             show_progress=args.progress_bar,
         )
 
-        prompt.submit(**kwargs)
-        negative_prompt.submit(**kwargs)
-        stable_diffusion.click(**kwargs)
+        prompt_submit = prompt.submit(**kwargs)
+        neg_prompt_submit = negative_prompt.submit(**kwargs)
+        generate_click = stable_diffusion.click(**kwargs)
+        clear_queue.click(
+            fn=None, cancels=[prompt_submit, neg_prompt_submit, generate_click]
+        )

apps/stable_diffusion/web/ui/txt2img_ui.py

@@ -1,7 +1,5 @@
-import os
-import sys
-import glob
 from pathlib import Path
+import os
 import gradio as gr
 from PIL import Image
 from apps.stable_diffusion.scripts import txt2img_inf
@@ -9,9 +7,12 @@ from apps.stable_diffusion.src import prompt_examples, args
 from apps.stable_diffusion.web.ui.utils import (
     available_devices,
     nodlogo_loc,
+    get_custom_model_path,
+    get_custom_model_files,
+    scheduler_list_txt2img,
+    predefined_models,
 )
 
-
 with gr.Blocks(title="Text-to-Image") as txt2img_web:
     with gr.Row(elem_id="ui_title"):
         nod_logo = Image.open(nodlogo_loc)
@@ -27,39 +28,33 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
         with gr.Row():
             with gr.Column(scale=1, min_width=600):
                 with gr.Row():
-                    ckpt_path = (
-                        Path(args.ckpt_dir)
-                        if args.ckpt_dir
-                        else Path(Path.cwd(), "models")
-                    )
-                    ckpt_path.mkdir(parents=True, exist_ok=True)
-                    types = (
-                        "*.ckpt",
-                        "*.safetensors",
-                    )  # the tuple of file types
-                    ckpt_files = ["None"]
-                    for extn in types:
-                        files = glob.glob(os.path.join(ckpt_path, extn))
-                        ckpt_files.extend(files)
-                    custom_model = gr.Dropdown(
-                        label=f"Models (Custom Model path: {ckpt_path})",
-                        value=args.ckpt_loc if args.ckpt_loc else "None",
-                        choices=ckpt_files
-                        + [
-                            "Linaqruf/anything-v3.0",
-                            "prompthero/openjourney",
-                            "wavymulder/Analog-Diffusion",
-                            "stabilityai/stable-diffusion-2-1",
-                            "stabilityai/stable-diffusion-2-1-base",
-                            "CompVis/stable-diffusion-v1-4",
-                        ],
-                    )
-                    hf_model_id = gr.Textbox(
-                        placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3",
-                        value="",
-                        label="HuggingFace Model ID",
-                        lines=3,
-                    )
+                    with gr.Column(scale=10):
+                        with gr.Row():
+                            custom_model = gr.Dropdown(
+                                label=f"Models (Custom Model path: {get_custom_model_path()})",
+                                elem_id="custom_model",
+                                value=os.path.basename(args.ckpt_loc)
+                                if args.ckpt_loc
+                                else "None",
+                                choices=["None"]
+                                + get_custom_model_files()
+                                + predefined_models,
+                            )
+                            hf_model_id = gr.Textbox(
+                                elem_id="hf_model_id",
+                                placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3",
+                                value="",
+                                label="HuggingFace Model ID",
+                                lines=3,
+                            )
+                    with gr.Column(scale=1, min_width=170):
+                        png_info_img = gr.Image(
+                            label="Import PNG info",
+                            elem_id="txt2img_prompt_image",
+                            type="pil",
+                            tool="None",
+                            visible=True,
+                        )
 
                 with gr.Group(elem_id="prompt_box_outer"):
                     prompt = gr.Textbox(
@@ -74,21 +69,30 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
                         lines=1,
                         elem_id="negative_prompt_box",
                     )
+                with gr.Accordion(
+                    label="Lora based inference option", open=False
+                ):
+                    with gr.Row():
+                        lora_weights = gr.Dropdown(
+                            label=f"Standlone LoRA weights (Path: {get_custom_model_path()})",
+                            elem_id="lora_weights",
+                            value="None",
+                            choices=["None"] + get_custom_model_files(),
+                        )
+                        lora_hf_id = gr.Textbox(
+                            elem_id="lora_hf_id",
+                            placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4",
+                            value="",
+                            label="HuggingFace Model ID",
+                            lines=3,
+                        )
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():
                         scheduler = gr.Dropdown(
+                            elem_id="scheduler",
                             label="Scheduler",
                             value=args.scheduler,
-                            choices=[
-                                "DDIM",
-                                "PNDM",
-                                "LMSDiscrete",
-                                "KDPM2Discrete",
-                                "DPMSolverMultistep",
-                                "EulerDiscrete",
-                                "EulerAncestralDiscrete",
-                                "SharkEulerDiscrete",
-                            ],
+                            choices=scheduler_list_txt2img,
                         )
                         with gr.Group():
                             save_metadata_to_png = gr.Checkbox(
@@ -159,19 +163,25 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
                         value=args.seed, precision=0, label="Seed"
                     )
                     device = gr.Dropdown(
+                        elem_id="device",
                         label="Device",
                         value=available_devices[0],
                         choices=available_devices,
                     )
                 with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        None,
-                        inputs=[],
-                        outputs=[seed],
-                        _js="() => Math.floor(Math.random() * 4294967295)",
-                    )
-                    stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=2):
+                        random_seed = gr.Button("Randomize Seed")
+                        random_seed.click(
+                            None,
+                            inputs=[],
+                            outputs=[seed],
+                            _js="() => -1",
+                        )
+                    with gr.Column(scale=6):
+                        stable_diffusion = gr.Button("Generate Image(s)")
+                    with gr.Column(scale=1, min_width=150):
+                        clear_queue = gr.Button("Clear Queue")
+
                 with gr.Accordion(label="Prompt Examples!", open=False):
                     ex = gr.Examples(
                         examples=prompt_examples,
@@ -226,11 +236,40 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
                 max_length,
                 save_metadata_to_json,
                 save_metadata_to_png,
+                lora_weights,
+                lora_hf_id,
             ],
             outputs=[txt2img_gallery, std_output],
             show_progress=args.progress_bar,
         )
 
-        prompt.submit(**kwargs)
-        negative_prompt.submit(**kwargs)
-        stable_diffusion.click(**kwargs)
+        prompt_submit = prompt.submit(**kwargs)
+        neg_prompt_submit = negative_prompt.submit(**kwargs)
+        generate_click = stable_diffusion.click(**kwargs)
+        clear_queue.click(
+            fn=None, cancels=[prompt_submit, neg_prompt_submit, generate_click]
+        )
+
+        from apps.stable_diffusion.web.utils.png_metadata import (
+            import_png_metadata,
+        )
+
+        png_info_img.change(
+            fn=import_png_metadata,
+            inputs=[
+                png_info_img,
+            ],
+            outputs=[
+                png_info_img,
+                prompt,
+                negative_prompt,
+                steps,
+                scheduler,
+                guidance_scale,
+                seed,
+                width,
+                height,
+                custom_model,
+                hf_model_id,
+            ],
+        )

apps/stable_diffusion/web/ui/utils.py

@@ -1,6 +1,44 @@
 import os
 import sys
 from apps.stable_diffusion.src import get_available_devices
+import glob
+from pathlib import Path
+from apps.stable_diffusion.src import args
+
+custom_model_filetypes = (
+    "*.ckpt",
+    "*.safetensors",
+)  # the tuple of file types
+
+scheduler_list = [
+    "DDIM",
+    "PNDM",
+    "DPMSolverMultistep",
+    "EulerAncestralDiscrete",
+]
+scheduler_list_txt2img = [
+    "DDIM",
+    "PNDM",
+    "LMSDiscrete",
+    "KDPM2Discrete",
+    "DPMSolverMultistep",
+    "EulerDiscrete",
+    "EulerAncestralDiscrete",
+    "SharkEulerDiscrete",
+]
+
+predefined_models = [
+    "Linaqruf/anything-v3.0",
+    "prompthero/openjourney",
+    "wavymulder/Analog-Diffusion",
+    "stabilityai/stable-diffusion-2-1",
+    "stabilityai/stable-diffusion-2-1-base",
+    "CompVis/stable-diffusion-v1-4",
+]
+predefined_paint_models = [
+    "runwayml/stable-diffusion-inpainting",
+    "stabilityai/stable-diffusion-2-inpainting",
+]
 
 
 def resource_path(relative_path):
@@ -11,5 +49,24 @@ def resource_path(relative_path):
     return os.path.join(base_path, relative_path)
 
 
+def get_custom_model_path():
+    return Path(args.ckpt_dir) if args.ckpt_dir else Path(Path.cwd(), "models")
+
+
+def get_custom_model_pathfile(custom_model_name):
+    return os.path.join(get_custom_model_path(), custom_model_name)
+
+
+def get_custom_model_files():
+    ckpt_files = []
+    for extn in custom_model_filetypes:
+        files = [
+            os.path.basename(x)
+            for x in glob.glob(os.path.join(get_custom_model_path(), extn))
+        ]
+        ckpt_files.extend(files)
+    return sorted(ckpt_files, key=str.casefold)
+
+
 nodlogo_loc = resource_path("logos/nod-logo.png")
 available_devices = get_available_devices()

apps/stable_diffusion/web/utils/png_metadata.py

@@ -0,0 +1,148 @@
+import re
+from pathlib import Path
+from apps.stable_diffusion.web.ui.txt2img_ui import (
+    png_info_img,
+    prompt,
+    negative_prompt,
+    steps,
+    scheduler,
+    guidance_scale,
+    seed,
+    width,
+    height,
+    custom_model,
+    hf_model_id,
+)
+from apps.stable_diffusion.web.ui.utils import (
+    get_custom_model_pathfile,
+    scheduler_list_txt2img,
+    predefined_models,
+)
+
+re_param_code = r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)'
+re_param = re.compile(re_param_code)
+re_imagesize = re.compile(r"^(\d+)x(\d+)$")
+
+
+def parse_generation_parameters(x: str):
+    res = {}
+    prompt = ""
+    negative_prompt = ""
+    done_with_prompt = False
+
+    *lines, lastline = x.strip().split("\n")
+    if len(re_param.findall(lastline)) < 3:
+        lines.append(lastline)
+        lastline = ""
+
+    for i, line in enumerate(lines):
+        line = line.strip()
+        if line.startswith("Negative prompt:"):
+            done_with_prompt = True
+            line = line[16:].strip()
+
+        if done_with_prompt:
+            negative_prompt += ("" if negative_prompt == "" else "\n") + line
+        else:
+            prompt += ("" if prompt == "" else "\n") + line
+
+    res["Prompt"] = prompt
+    res["Negative prompt"] = negative_prompt
+
+    for k, v in re_param.findall(lastline):
+        v = v[1:-1] if v[0] == '"' and v[-1] == '"' else v
+        m = re_imagesize.match(v)
+        if m is not None:
+            res[k + "-1"] = m.group(1)
+            res[k + "-2"] = m.group(2)
+        else:
+            res[k] = v
+
+    # Missing CLIP skip means it was set to 1 (the default)
+    if "Clip skip" not in res:
+        res["Clip skip"] = "1"
+
+    hypernet = res.get("Hypernet", None)
+    if hypernet is not None:
+        res[
+            "Prompt"
+        ] += f"""<hypernet:{hypernet}:{res.get("Hypernet strength", "1.0")}>"""
+
+    if "Hires resize-1" not in res:
+        res["Hires resize-1"] = 0
+        res["Hires resize-2"] = 0
+
+    return res
+
+
+def import_png_metadata(pil_data):
+    try:
+        png_info = pil_data.info["parameters"]
+        metadata = parse_generation_parameters(png_info)
+        png_hf_model_id = ""
+        png_custom_model = ""
+
+        if "Model" in metadata:
+            # Remove extension from model info
+            if metadata["Model"].endswith(".safetensors") or metadata[
+                "Model"
+            ].endswith(".ckpt"):
+                metadata["Model"] = Path(metadata["Model"]).stem
+            # Check for the model name match with one of the local ckpt or safetensors files
+            if Path(
+                get_custom_model_pathfile(metadata["Model"] + ".ckpt")
+            ).is_file():
+                png_custom_model = metadata["Model"] + ".ckpt"
+            if Path(
+                get_custom_model_pathfile(metadata["Model"] + ".safetensors")
+            ).is_file():
+                png_custom_model = metadata["Model"] + ".safetensors"
+            # Check for a model match with one of the default model list (ex: "Linaqruf/anything-v3.0")
+            if metadata["Model"] in predefined_models:
+                png_custom_model = metadata["Model"]
+            # If nothing had matched, check vendor/hf_model_id
+            if not png_custom_model and metadata["Model"].count("/"):
+                png_hf_model_id = metadata["Model"]
+            # No matching model was found
+            if not png_custom_model and not png_hf_model_id:
+                print(
+                    "Import PNG info: Unable to find a matching model for %s"
+                    % metadata["Model"]
+                )
+
+        outputs = {
+            png_info_img: None,
+            negative_prompt: metadata["Negative prompt"],
+            steps: int(metadata["Steps"]),
+            guidance_scale: float(metadata["CFG scale"]),
+            seed: int(metadata["Seed"]),
+            width: float(metadata["Size-1"]),
+            height: float(metadata["Size-2"]),
+        }
+        if "Model" in metadata and png_custom_model:
+            outputs[custom_model] = png_custom_model
+            outputs[hf_model_id] = ""
+        if "Model" in metadata and png_hf_model_id:
+            outputs[custom_model] = "None"
+            outputs[hf_model_id] = png_hf_model_id
+        if "Prompt" in metadata:
+            outputs[prompt] = metadata["Prompt"]
+        if "Sampler" in metadata:
+            if metadata["Sampler"] in scheduler_list_txt2img:
+                outputs[scheduler] = metadata["Sampler"]
+            else:
+                print(
+                    "Import PNG info: Unable to find a scheduler for %s"
+                    % metadata["Sampler"]
+                )
+
+        return outputs
+
+    except Exception as ex:
+        if pil_data and pil_data.info.get("parameters"):
+            print("import_png_metadata failed with %s" % ex)
+        pass
+
+    return {
+        png_info_img: None,
+    }

cpp/README.md

@@ -40,7 +40,7 @@ cmake --build build/
 *Prepare the model*
 ```bash
 wget https://storage.googleapis.com/shark_tank/latest/resnet50_tf/resnet50_tf.mlir
-iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --iree-llvm-embedded-linker-path=`python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`/iree/compiler/tools/../_mlir_libs/iree-lld --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --mlir-pass-pipeline-crash-reproducer=ist/core-reproducer.mlir --iree-llvm-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 resnet50_tf.mlir -o resnet50_tf.vmfb
+iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --iree-llvmcpu-embedded-linker-path=`python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`/iree/compiler/tools/../_mlir_libs/iree-lld --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --mlir-pass-pipeline-crash-reproducer=ist/core-reproducer.mlir --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 resnet50_tf.mlir -o resnet50_tf.vmfb
 ```
 *Prepare the input*
 
@@ -65,18 +65,18 @@ A tool for benchmarking other models is built and can be invoked with a command
 see `./build/vulkan_gui/iree-vulkan-gui --help` for an explanation on the function input. For example, stable diffusion unet can be tested with the following commands:
 ```bash
 wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/stable_diff_tf.mlir
-iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvm-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 stable_diff_tf.mlir -o stable_diff_tf.vmfb
+iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 stable_diff_tf.mlir -o stable_diff_tf.vmfb
 ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=2x4x64x64xf32 --function_input=1xf32 --function_input=2x77x768xf32
 ```
 VAE and Autoencoder are also available
 ```bash
 # VAE
 wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/vae_tf/vae.mlir
-iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvm-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 vae.mlir -o vae.vmfb
+iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 vae.mlir -o vae.vmfb
 ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=1x4x64x64xf32
 
 # CLIP Autoencoder
 wget https://storage.googleapis.com/shark_tank/quinn/stable_diff_tf/clip_tf/clip_autoencoder.mlir
-iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvm-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 clip_autoencoder.mlir -o clip_autoencoder.vmfb
+iree-compile --iree-input-type=mhlo --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host -iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 clip_autoencoder.mlir -o clip_autoencoder.vmfb
 ./build/vulkan_gui/iree-vulkan-gui --module-file=stable_diff_tf.vmfb --function_input=1x77xi32 --function_input=1x77xi32
 ```

generate_sharktank.py

@@ -52,7 +52,7 @@ def save_torch_model(torch_model_list):
 
             tracing_required = False if tracing_required == "False" else True
             is_dynamic = False if is_dynamic == "False" else True
-
+            print("generating artifacts for: " + torch_model_name)
             model = None
             input = None
             if model_type == "stable_diffusion":
@@ -105,12 +105,6 @@ def save_torch_model(torch_model_list):
                 dir=torch_model_dir,
                 model_name=torch_model_name,
             )
-            mlir_hash = create_hash(
-                os.path.join(
-                    torch_model_dir, torch_model_name + "_torch" + ".mlir"
-                )
-            )
-            np.save(os.path.join(torch_model_dir, "hash"), np.array(mlir_hash))
             # Generate torch dynamic models.
             if is_dynamic:
                 mlir_importer.import_debug(
@@ -276,6 +270,9 @@ if __name__ == "__main__":
         os.path.dirname(__file__), "tank", "tflite", "tflite_model_list.csv"
     )
 
+    save_torch_model(
+        os.path.join(os.path.dirname(__file__), "tank", "torch_sd_list.csv")
+    )
     save_torch_model(torch_model_csv)
     save_tf_model(tf_model_csv)
     save_tflite_model(tflite_model_csv)

process_skipfiles.py

@@ -7,14 +7,35 @@ import fileinput
 from pathlib import Path
 
 # Diffusers 0.13.1 fails with transformers __init.py errros in BLIP. So remove it for now until we fork it
-pix2pix_file = Path(
-    get_python_lib()
-    + "/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py"
+pix2pix_init = Path(get_python_lib() + "/diffusers/__init__.py")
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "Pix2Pix" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
+pix2pix_init = Path(get_python_lib() + "/diffusers/pipelines/__init__.py")
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "Pix2Pix" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
+pix2pix_init = Path(
+    get_python_lib() + "/diffusers/pipelines/stable_diffusion/__init__.py"
 )
-if pix2pix_file.exists():
-    print("Removing..%s", pix2pix_file)
-    pix2pix_file.unlink()
-
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "StableDiffusionPix2PixZeroPipeline" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
 
 path_to_skipfiles = Path(get_python_lib() + "/torch/_dynamo/skipfiles.py")
 
@@ -42,3 +63,16 @@ for line in fileinput.input(path_to_skipfiles, inplace=True):
                 print(line, end="")
     else:
         print(line, end="")
+
+# For getting around scikit-image's packaging, laze_loader has had a patch merged but yet to be released.
+# Refer: https://github.com/scientific-python/lazy_loader
+path_to_lazy_loader = Path(get_python_lib() + "/lazy_loader/__init__.py")
+
+for line in fileinput.input(path_to_lazy_loader, inplace=True):
+    if 'stubfile = filename if filename.endswith("i")' in line:
+        print(
+            '    stubfile = (filename if filename.endswith("i") else f"{os.path.splitext(filename)[0]}.pyi")',
+            end="",
+        )
+    else:
+        print(line, end="")

pyproject.toml

@@ -10,3 +10,8 @@ requires = [
     "iree-runtime>=20221022.190",
 ]
 build-backend = "setuptools.build_meta"
+
+[tool.black]
+line-length = 79
+include = '\.pyi?$'
+

requirements.txt

@@ -16,7 +16,7 @@ parameterized
 
 # Add transformers, diffusers and scipy since it most commonly used
 transformers
-diffusers @ git+https://github.com/nod-ai/diffusers@stable_stencil
+diffusers
 scipy
 ftfy
 gradio
@@ -24,6 +24,7 @@ altair
 omegaconf
 safetensors
 opencv-python
+scikit-image
 
 # Keep PyInstaller at the end. Sometimes Windows Defender flags it but most folks can continue even if it errors
 pefile

setup_venv.ps1

@@ -89,7 +89,7 @@ else {python -m venv .\shark.venv\}
 python -m pip install --upgrade pip
 pip install wheel
 pip install -r requirements.txt
-pip install --pre torch-mlir==20230228.763 torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://llvm.github.io/torch-mlir/package-index/
+pip install --pre torch-mlir torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://llvm.github.io/torch-mlir/package-index/
 pip install --upgrade -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html iree-compiler iree-runtime
 Write-Host "Building SHARK..."
 pip install -e . -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html

setup_venv.sh

@@ -78,9 +78,9 @@ $PYTHON -m pip install --upgrade -r "$TD/requirements.txt"
 if [ "$torch_mlir_bin" = true ]; then
   if [[ $(uname -s) = 'Darwin' ]]; then
     echo "MacOS detected. Installing torch-mlir from .whl, to avoid dependency problems with torch."
-    $PYTHON -m pip install --pre --no-cache-dir torch-mlir==20230228.763 -f https://llvm.github.io/torch-mlir/package-index/ -f https://download.pytorch.org/whl/nightly/torch/
+    $PYTHON -m pip install --pre --no-cache-dir torch-mlir -f https://llvm.github.io/torch-mlir/package-index/ -f https://download.pytorch.org/whl/nightly/torch/
   else
-    $PYTHON -m pip install --pre torch-mlir==20230228.763 -f https://llvm.github.io/torch-mlir/package-index/
+    $PYTHON -m pip install --pre torch-mlir -f https://llvm.github.io/torch-mlir/package-index/
     if [ $? -eq 0 ];then
       echo "Successfully Installed torch-mlir"
     else

shark/examples/shark_inference/llama/README.md

@@ -0,0 +1,18 @@
+# SHARK LLaMA
+
+## TORCH-MLIR Version
+
+```
+https://github.com/nod-ai/torch-mlir.git
+```
+Then check out the `complex` branch and `git submodule update --init` and then build with `.\build_tools\python_deploy\build_windows.ps1`
+
+### Setup & Run
+```
+git clone https://github.com/nod-ai/llama.git
+```
+Then in this repository
+```
+pip install -e .
+python llama/shark_model.py
+```

shark/examples/shark_inference/sharded_bloom.py

@@ -7,9 +7,13 @@
 # -de --device: the device you want to run bloom on.  E.G. cpu, cuda
 # -c, --recompile: set to true if you want to recompile to vmfb.
 # -d, --download: set to true if you want to redownload the mlir files
+# -cm, --create_mlirs: set to true if you want to create the mlir files from scratch.  please make sure you have transformers 4.21.2 before using this option
 # -t --token_count: the number of tokens you want to generate
 # -pr --prompt: the prompt you want to feed to the model
-# -m --model_namme: the name of the model, e.g. bloom-560m
+# -m --model_name: the name of the model, e.g. bloom-560m
+#
+# If you don't specify a prompt when you run this example, you will be able to give prompts through the terminal.  Run the
+# example in this way if you want to run multiple examples without reinitializing the model
 #####################################################################################
 
 import os
@@ -26,6 +30,7 @@ import sys
 import argparse
 import json
 import urllib.request
+import subprocess
 
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch._decomp import get_decompositions
@@ -306,31 +311,8 @@ def _prepare_attn_mask(
 
 def download_model(destination_folder, model_name):
     download_public_file(
-        f"https://{model_name}/config.json", destination_folder
+        f"gs://shark_tank/sharded_bloom/{model_name}/", destination_folder
     )
-    f = open(f"{destination_folder}/config.json")
-    config = json.load(f)
-    f.close()
-    n_blocks = config["n_layer"]
-    download_public_file(
-        f"https://{model_name}/lm_head.mlir", destination_folder
-    )
-    download_public_file(f"https://{model_name}/ln_f.mlir", destination_folder)
-    download_public_file(
-        f"https://{model_name}/word_embeddings.mlir", destination_folder
-    )
-    download_public_file(
-        f"https://{model_name}/word_embeddings_layernorm.mlir",
-        destination_folder,
-    )
-    download_public_file(
-        f"https://{model_name}/tokenizer.json", destination_folder
-    )
-
-    for i in range(n_blocks):
-        download_public_file(
-            f"https://{model_name}/bloom_block_{i}.mlir", destination_folder
-        )
 
 
 def compile_embeddings(embeddings_layer, input_ids, path):
@@ -653,6 +635,75 @@ def create_mlirs(destination_folder, model_name):
     )
 
 
+def run_large_model(
+    token_count,
+    recompile,
+    model_path,
+    prompt,
+    device_list,
+    script_path,
+    device,
+):
+    f = open(f"{model_path}/prompt.txt", "w+")
+    f.write(prompt)
+    f.close()
+    for i in range(token_count):
+        if i == 0:
+            will_compile = recompile
+        else:
+            will_compile = False
+            f = open(f"{model_path}/prompt.txt", "r")
+            prompt = f.read()
+            f.close()
+
+        subprocess.run(
+            [
+                "python",
+                script_path,
+                model_path,
+                "start",
+                str(will_compile),
+                "cpu",
+                "None",
+                prompt,
+            ]
+        )
+        for i in range(config["n_layer"]):
+            if device_list is not None:
+                device_idx = str(device_list[i % len(device_list)])
+            else:
+                device_idx = "None"
+            subprocess.run(
+                [
+                    "python",
+                    script_path,
+                    model_path,
+                    str(i),
+                    str(will_compile),
+                    device,
+                    device_idx,
+                    prompt,
+                ]
+            )
+        subprocess.run(
+            [
+                "python",
+                script_path,
+                model_path,
+                "end",
+                str(will_compile),
+                "cpu",
+                "None",
+                prompt,
+            ]
+        )
+
+    f = open(f"{model_path}/prompt.txt", "r")
+    output = f.read()
+    f.close()
+    print(output)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(prog="Bloom-560m")
     parser.add_argument("-p", "--model_path")
@@ -662,37 +713,130 @@ if __name__ == "__main__":
     parser.add_argument("-d", "--download", default=False, type=bool)
     parser.add_argument("-t", "--token_count", default=10, type=int)
     parser.add_argument("-m", "--model_name", default="bloom-560m")
+    parser.add_argument("-cm", "--create_mlirs", default=False, type=bool)
+
+    parser.add_argument(
+        "-lm", "--large_model_memory_efficient", default=False, type=bool
+    )
+
     parser.add_argument(
         "-pr",
         "--prompt",
-        default="The SQL command to extract all the users whose name starts with A is: ",
+        default=None,
     )
     args = parser.parse_args()
 
+    if args.create_mlirs and args.large_model_memory_efficient:
+        print(
+            "Warning: If you need to use memory efficient mode, you probably want to use 'download' instead"
+        )
+
+    if not os.path.isdir(args.model_path):
+        os.mkdir(args.model_path)
+
     if args.device_list is not None:
         args.device_list = json.loads(args.device_list)
 
     if args.device == "cuda" and args.device_list is not None:
         IS_CUDA = True
         from cuda.cudart import cudaSetDevice
+    if args.download and args.create_mlirs:
+        print(
+            "WARNING: It is not advised to turn on both download and create_mlirs"
+        )
     if args.download:
-        # download_model(args.model_path, args.model_name)
+        download_model(args.model_path, args.model_name)
+    if args.create_mlirs:
         create_mlirs(args.model_path, args.model_name)
     from transformers import AutoTokenizer, AutoModelForCausalLM, BloomConfig
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_path)
-    input_ids = tokenizer.encode(args.prompt, return_tensors="pt")
+    if args.prompt is not None:
+        input_ids = tokenizer.encode(args.prompt, return_tensors="pt")
 
-    shardedbloom = ShardedBloom(args.model_path)
-    shardedbloom.init_layers(
-        device=args.device, replace=args.recompile, device_idx=args.device_list
-    )
-    shardedbloom.load_layers()
+    if args.large_model_memory_efficient:
+        f = open(f"{args.model_path}/config.json")
+        config = json.load(f)
+        f.close()
 
-    for _ in range(args.token_count):
-        next_token = shardedbloom.forward_pass(
-            torch.tensor(input_ids), device=args.device
+        self_path = os.path.dirname(os.path.abspath(__file__))
+        script_path = os.path.join(self_path, "sharded_bloom_large_models.py")
+
+        if args.prompt is not None:
+            run_large_model(
+                args.token_count,
+                args.recompile,
+                args.model_path,
+                args.prompt,
+                args.device_list,
+                script_path,
+                args.device,
+            )
+
+        else:
+            while True:
+                prompt = input("Enter Prompt: ")
+                try:
+                    token_count = int(
+                        input("Enter number of tokens you want to generate: ")
+                    )
+                except:
+                    print(
+                        "Invalid integer entered.  Using default value of 10"
+                    )
+                    token_count = 10
+
+                run_large_model(
+                    token_count,
+                    args.recompile,
+                    args.model_path,
+                    prompt,
+                    args.device_list,
+                    script_path,
+                    args.device,
+                )
+
+    else:
+        shardedbloom = ShardedBloom(args.model_path)
+        shardedbloom.init_layers(
+            device=args.device,
+            replace=args.recompile,
+            device_idx=args.device_list,
         )
-        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)
+        shardedbloom.load_layers()
 
-    print(tokenizer.decode(input_ids.squeeze()))
+        if args.prompt is not None:
+            for _ in range(args.token_count):
+                next_token = shardedbloom.forward_pass(
+                    torch.tensor(input_ids), device=args.device
+                )
+                input_ids = torch.cat(
+                    [input_ids, next_token.unsqueeze(-1)], dim=-1
+                )
+
+            print(tokenizer.decode(input_ids.squeeze()))
+
+        else:
+            while True:
+                prompt = input("Enter Prompt: ")
+                try:
+                    token_count = int(
+                        input("Enter number of tokens you want to generate: ")
+                    )
+                except:
+                    print(
+                        "Invalid integer entered.  Using default value of 10"
+                    )
+                    token_count = 10
+
+                input_ids = tokenizer.encode(prompt, return_tensors="pt")
+
+                for _ in range(token_count):
+                    next_token = shardedbloom.forward_pass(
+                        torch.tensor(input_ids), device=args.device
+                    )
+                    input_ids = torch.cat(
+                        [input_ids, next_token.unsqueeze(-1)], dim=-1
+                    )
+
+                print(tokenizer.decode(input_ids.squeeze()))

shark/examples/shark_inference/sharded_bloom_large_models.py

@@ -0,0 +1,381 @@
+import sys
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, BloomConfig
+import re
+from shark.shark_inference import SharkInference
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+from transformers.models.bloom.modeling_bloom import (
+    BloomBlock,
+    build_alibi_tensor,
+)
+import time
+import json
+
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    batch_size, source_length = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else source_length
+
+    expanded_mask = (
+        mask[:, None, None, :]
+        .expand(batch_size, 1, tgt_len, source_length)
+        .to(dtype)
+    )
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(
+        inverted_mask.to(torch.bool), torch.finfo(dtype).min
+    )
+
+
+def _prepare_attn_mask(
+    attention_mask, input_shape, inputs_embeds, past_key_values_length
+):
+    # create causal mask
+    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+    combined_attention_mask = None
+    if input_shape[-1] > 1:
+        combined_attention_mask = _make_causal_mask(
+            input_shape,
+            inputs_embeds.dtype,
+            past_key_values_length=past_key_values_length,
+        ).to(attention_mask.device)
+
+    if attention_mask is not None:
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        expanded_attn_mask = _expand_mask(
+            attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+        )
+        combined_attention_mask = (
+            expanded_attn_mask
+            if combined_attention_mask is None
+            else expanded_attn_mask + combined_attention_mask
+        )
+
+    return combined_attention_mask
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    past_key_values_length: int = 0,
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    batch_size, target_length = input_ids_shape
+    mask = torch.full((target_length, target_length), torch.finfo(dtype).min)
+    mask_cond = torch.arange(mask.size(-1))
+    intermediate_mask = mask_cond < (mask_cond + 1).view(mask.size(-1), 1)
+    mask.masked_fill_(intermediate_mask, 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat(
+            [
+                torch.zeros(
+                    target_length, past_key_values_length, dtype=dtype
+                ),
+                mask,
+            ],
+            dim=-1,
+        )
+    expanded_mask = mask[None, None, :, :].expand(
+        batch_size, 1, target_length, target_length + past_key_values_length
+    )
+    return expanded_mask
+
+
+if __name__ == "__main__":
+    working_dir = sys.argv[1]
+    layer_name = sys.argv[2]
+    will_compile = sys.argv[3]
+    device = sys.argv[4]
+    device_idx = sys.argv[5]
+    prompt = sys.argv[6]
+
+    if device_idx.lower().strip() == "none":
+        device_idx = None
+    else:
+        device_idx = int(device_idx)
+
+    if will_compile.lower().strip() == "true":
+        will_compile = True
+    else:
+        will_compile = False
+
+    f = open(f"{working_dir}/config.json")
+    config = json.load(f)
+    f.close()
+
+    layers_initialized = False
+    try:
+        n_embed = config["n_embed"]
+    except KeyError:
+        n_embed = config["hidden_size"]
+    vocab_size = config["vocab_size"]
+    n_layer = config["n_layer"]
+    try:
+        n_head = config["num_attention_heads"]
+    except KeyError:
+        n_head = config["n_head"]
+
+    if not os.path.isdir(working_dir):
+        os.mkdir(working_dir)
+
+    if layer_name == "start":
+        tokenizer = AutoTokenizer.from_pretrained(working_dir)
+        input_ids = tokenizer.encode(prompt, return_tensors="pt")
+
+        mlir_str = ""
+
+        if will_compile:
+            f = open(f"{working_dir}/word_embeddings.mlir", encoding="utf-8")
+            mlir_str = f.read()
+            f.close()
+
+            mlir_str = bytes(mlir_str, "utf-8")
+
+        shark_module = SharkInference(
+            mlir_str,
+            device="cpu",
+            mlir_dialect="tm_tensor",
+            device_idx=None,
+        )
+
+        if will_compile:
+            shark_module.save_module(
+                module_name=f"{working_dir}/word_embeddings",
+                extra_args=[
+                    "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                    "--iree-stream-resource-max-allocation-size=1000000000",
+                    "--iree-codegen-check-ir-before-llvm-conversion=false",
+                ],
+            )
+
+        shark_module.load_module(f"{working_dir}/word_embeddings.vmfb")
+        input_embeds = shark_module(
+            inputs=(input_ids,), function_name="forward"
+        )
+        input_embeds = torch.tensor(input_embeds).float()
+
+        mlir_str = ""
+
+        if will_compile:
+            f = open(
+                f"{working_dir}/word_embeddings_layernorm.mlir",
+                encoding="utf-8",
+            )
+            mlir_str = f.read()
+            f.close()
+
+        shark_module = SharkInference(
+            mlir_str,
+            device="cpu",
+            mlir_dialect="tm_tensor",
+            device_idx=None,
+        )
+
+        if will_compile:
+            shark_module.save_module(
+                module_name=f"{working_dir}/word_embeddings_layernorm",
+                extra_args=[
+                    "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                    "--iree-stream-resource-max-allocation-size=1000000000",
+                    "--iree-codegen-check-ir-before-llvm-conversion=false",
+                ],
+            )
+
+        shark_module.load_module(
+            f"{working_dir}/word_embeddings_layernorm.vmfb"
+        )
+        hidden_states = shark_module(
+            inputs=(input_embeds,), function_name="forward"
+        )
+        hidden_states = torch.tensor(hidden_states).float()
+
+        torch.save(hidden_states, f"{working_dir}/hidden_states_0.pt")
+
+        attention_mask = torch.ones(
+            [hidden_states.shape[0], len(input_ids[0])]
+        )
+
+        attention_mask = torch.tensor(attention_mask).float()
+
+        alibi = build_alibi_tensor(
+            attention_mask,
+            n_head,
+            hidden_states.dtype,
+            device="cpu",
+        )
+
+        torch.save(alibi, f"{working_dir}/alibi.pt")
+
+        causal_mask = _prepare_attn_mask(
+            attention_mask, input_ids.size(), input_embeds, 0
+        )
+        causal_mask = torch.tensor(causal_mask).float()
+
+        torch.save(causal_mask, f"{working_dir}/causal_mask.pt")
+
+    elif layer_name in [str(x) for x in range(n_layer)]:
+        hidden_states = torch.load(
+            f"{working_dir}/hidden_states_{layer_name}.pt"
+        )
+        alibi = torch.load(f"{working_dir}/alibi.pt")
+        causal_mask = torch.load(f"{working_dir}/causal_mask.pt")
+
+        mlir_str = ""
+
+        if will_compile:
+            f = open(
+                f"{working_dir}/bloom_block_{layer_name}.mlir",
+                encoding="utf-8",
+            )
+            mlir_str = f.read()
+            f.close()
+
+            mlir_str = bytes(mlir_str, "utf-8")
+
+        shark_module = SharkInference(
+            mlir_str,
+            device=device,
+            mlir_dialect="tm_tensor",
+            device_idx=device_idx,
+        )
+
+        if will_compile:
+            shark_module.save_module(
+                module_name=f"{working_dir}/bloom_block_{layer_name}",
+                extra_args=[
+                    "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                    "--iree-stream-resource-max-allocation-size=1000000000",
+                    "--iree-codegen-check-ir-before-llvm-conversion=false",
+                ],
+            )
+
+        shark_module.load_module(
+            f"{working_dir}/bloom_block_{layer_name}.vmfb"
+        )
+
+        output = shark_module(
+            inputs=(
+                hidden_states.detach().numpy(),
+                alibi.detach().numpy(),
+                causal_mask.detach().numpy(),
+            ),
+            function_name="forward",
+        )
+
+        hidden_states = torch.tensor(output[0]).float()
+
+        torch.save(
+            hidden_states,
+            f"{working_dir}/hidden_states_{int(layer_name) + 1}.pt",
+        )
+
+    elif layer_name == "end":
+        mlir_str = ""
+
+        if will_compile:
+            f = open(f"{working_dir}/ln_f.mlir", encoding="utf-8")
+            mlir_str = f.read()
+            f.close()
+
+            mlir_str = bytes(mlir_str, "utf-8")
+
+        shark_module = SharkInference(
+            mlir_str,
+            device="cpu",
+            mlir_dialect="tm_tensor",
+            device_idx=None,
+        )
+
+        if will_compile:
+            shark_module.save_module(
+                module_name=f"{working_dir}/ln_f",
+                extra_args=[
+                    "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                    "--iree-stream-resource-max-allocation-size=1000000000",
+                    "--iree-codegen-check-ir-before-llvm-conversion=false",
+                ],
+            )
+
+        shark_module.load_module(f"{working_dir}/ln_f.vmfb")
+
+        hidden_states = torch.load(f"{working_dir}/hidden_states_{n_layer}.pt")
+
+        hidden_states = shark_module(
+            inputs=(hidden_states,), function_name="forward"
+        )
+
+        mlir_str = ""
+
+        if will_compile:
+            f = open(f"{working_dir}/lm_head.mlir", encoding="utf-8")
+            mlir_str = f.read()
+            f.close()
+
+            mlir_str = bytes(mlir_str, "utf-8")
+
+        if config["n_embed"] == 14336:
+
+            def get_state_dict():
+                d = torch.load(
+                    f"{working_dir}/pytorch_model_00001-of-00072.bin"
+                )
+                return OrderedDict(
+                    (k.replace("word_embeddings.", ""), v)
+                    for k, v in d.items()
+                )
+
+            def load_causal_lm_head():
+                linear = nn.utils.skip_init(
+                    nn.Linear, 14336, 250880, bias=False, dtype=torch.float
+                )
+                linear.load_state_dict(get_state_dict(), strict=False)
+                return linear.float()
+
+            lm_head = load_causal_lm_head()
+
+            logits = lm_head(torch.tensor(hidden_states).float())
+
+        else:
+            shark_module = SharkInference(
+                mlir_str,
+                device="cpu",
+                mlir_dialect="tm_tensor",
+                device_idx=None,
+            )
+
+            if will_compile:
+                shark_module.save_module(
+                    module_name=f"{working_dir}/lm_head",
+                    extra_args=[
+                        "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                        "--iree-stream-resource-max-allocation-size=1000000000",
+                        "--iree-codegen-check-ir-before-llvm-conversion=false",
+                    ],
+                )
+
+            shark_module.load_module(f"{working_dir}/lm_head.vmfb")
+
+            logits = shark_module(
+                inputs=(hidden_states,), function_name="forward"
+            )
+
+        logits = torch.tensor(logits).float()
+
+        tokenizer = AutoTokenizer.from_pretrained(working_dir)
+
+        next_token = tokenizer.decode(torch.argmax(logits[:, -1, :], dim=-1))
+
+        f = open(f"{working_dir}/prompt.txt", "w+")
+        f.write(prompt + next_token)
+        f.close()

shark/examples/shark_training/stable_diffusion/README.md

@@ -1,9 +1,6 @@
 # Stable Diffusion Fine Tuning
 
-## Installation
-
-<details>
-  <summary>Installation (Linux)</summary>
+## Installation (Linux)
 
 ### Activate shark.venv Virtual Environment
 
@@ -14,21 +11,21 @@ source shark.venv/bin/activate
 python -m pip install --upgrade pip
 ```
 
-### Install dependencies
+## Install dependencies
 
-# Run the following installation commands:
+### Run the following installation commands:
 ```
 pip install -U git+https://github.com/huggingface/diffusers.git
 pip install accelerate transformers ftfy
 ```
 
-# Build torch-mlir with the following branch:
+### Build torch-mlir with the following branch:
 
 Please cherry-pick this branch of torch-mlir: https://github.com/vivekkhandelwal1/torch-mlir/tree/sd-ops
 and build it locally. You can find the instructions for using locally build Torch-MLIR,
 here: https://github.com/nod-ai/SHARK#how-to-use-your-locally-built-iree--torch-mlir-with-shark
 
-### Run the Stable diffusion fine tuning
+## Run the Stable diffusion fine tuning
 
 To run the model with the default set of images and params, run:
 ```shell
@@ -42,3 +39,5 @@ For example, you can run the training for a limited set of steps via the dynamo
 ```
 python stable_diffusion_fine_tuning.py --training_steps=1 --inference_steps=1 --num_inference_samples=1 --train_batch_size=1 --use_torchdynamo=True
 ```
+
+You can also specify the device to be used via the flag `--device`. The default value is `cpu`, for GPU execution you can specify `--device="cuda"`.

shark/examples/shark_training/stable_diffusion/stable_diffusion_fine_tuning.py

@@ -445,8 +445,11 @@ freeze_params(params_to_freeze)
 
 
 # Move vae and unet to device
-vae.to(args.device)
-unet.to(args.device)
+# For the dynamo path default compilation device is `cpu`, since torch-mlir
+# supports only that. Therefore, convert to device only for PyTorch path.
+if not args.use_torchdynamo:
+    vae.to(args.device)
+    unet.to(args.device)
 
 # Keep vae in eval mode as we don't train it
 vae.eval()
@@ -663,7 +666,7 @@ def refbackend_torchdynamo_backend(
 def predictions(torch_func, jit_func, batchA, batchB):
     res = jit_func(batchA.numpy(), batchB.numpy())
     if res is not None:
-        prediction = torch.from_numpy(res)
+        prediction = res
     else:
         prediction = None
     return prediction
@@ -802,9 +805,13 @@ def training_function():
     progress_bar.set_description("Steps")
     global_step = 0
 
-    params__ = [i for i in text_encoder.get_input_embeddings().parameters()]
+    params_ = [i for i in text_encoder.get_input_embeddings().parameters()]
+    if args.use_torchdynamo:
+        print("******** TRAINING STARTED - TORCHYDNAMO PATH ********")
+    else:
+        print("******** TRAINING STARTED - PYTORCH PATH ********")
     print("Initial weights:")
-    print(params__, params__[0].shape)
+    print(params_, params_[0].shape)
 
     for epoch in range(num_train_epochs):
         text_encoder.train()
@@ -849,6 +856,7 @@ def training_function():
 
     # Create the pipeline using using the trained modules and save it.
     params__ = [i for i in text_encoder.get_input_embeddings().parameters()]
+    print("******** TRAINING PROCESS FINISHED ********")
     print("Updated weights:")
     print(params__, params__[0].shape)
     pipeline = StableDiffusionPipeline.from_pretrained(
@@ -880,7 +888,9 @@ pipe = StableDiffusionPipeline.from_pretrained(
     scheduler=DPMSolverMultistepScheduler.from_pretrained(
         hyperparameters["output_dir"], subfolder="scheduler"
     ),
-).to(args.device)
+)
+if not args.use_torchdynamo:
+    pipe.to(args.device)
 
 # Run the Stable Diffusion pipeline
 # Don't forget to use the placeholder token in your prompt

shark/iree_utils/_common.py

@@ -19,10 +19,14 @@ import sys
 import subprocess
 
 
-def run_cmd(cmd):
+def run_cmd(cmd, debug=False):
     """
     Inputs: cli command string.
     """
+    if debug:
+        print("IREE run command: \n\n")
+        print(cmd)
+        print("\n\n")
     try:
         result = subprocess.run(
             cmd,

shark/iree_utils/compile_utils.py

@@ -53,10 +53,10 @@ def get_iree_device_args(device, extra_args=[]):
 # Get the iree-compiler arguments given frontend.
 def get_iree_frontend_args(frontend):
     if frontend in ["torch", "pytorch", "linalg"]:
-        return ["--iree-llvm-target-cpu-features=host"]
+        return ["--iree-llvmcpu-target-cpu-features=host"]
     elif frontend in ["tensorflow", "tf", "mhlo"]:
         return [
-            "--iree-llvm-target-cpu-features=host",
+            "--iree-llvmcpu-target-cpu-features=host",
             "--iree-mhlo-demote-i64-to-i32=false",
             "--iree-flow-demote-i64-to-i32",
         ]
@@ -293,7 +293,8 @@ def get_iree_module(flatbuffer_blob, device, device_idx=None):
         haldriver = ireert.get_driver(device)
 
         haldevice = haldriver.create_device(
-            haldriver.query_available_devices()[device_idx]["device_id"]
+            haldriver.query_available_devices()[device_idx]["device_id"],
+            allocators=shark_args.device_allocator,
         )
         config = ireert.Config(device=haldevice)
     else:
@@ -402,5 +403,10 @@ def get_results(
 
 def get_iree_runtime_config(device):
     device = iree_device_map(device)
-    config = ireert.Config(device=ireert.get_device(device))
+    haldriver = ireert.get_driver(device)
+    haldevice = haldriver.create_device_by_uri(
+        device,
+        allocators=shark_args.device_allocator,
+    )
+    config = ireert.Config(device=haldevice)
     return config

shark/iree_utils/cpu_utils.py

@@ -44,4 +44,4 @@ def get_iree_cpu_args():
         error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)"
         raise Exception(error_message)
     print(f"Target triple found:{target_triple}")
-    return [f"-iree-llvm-target-triple={target_triple}"]
+    return [f"--iree-llvmcpu-target-triple={target_triple}"]

shark/parser.py

@@ -108,4 +108,14 @@ parser.add_argument(
     help="Enables the --iree-flow-enable-conv-winograd-transform flag.",
 )
 
+parser.add_argument(
+    "--device_allocator",
+    type=str,
+    nargs="*",
+    default=[],
+    help="Specifies one or more HAL device allocator specs "
+    "to augment the base device allocator",
+    choices=["debug", "caching"],
+)
+
 shark_args, unknown = parser.parse_known_args()

shark/shark_downloader.py

@@ -150,10 +150,14 @@ def download_model(
     if not check_dir_exists(
         model_dir_name, frontend=frontend, dynamic=dyn_str
     ):
-        print(f"Downloading artifacts for model {model_name}...")
+        print(
+            f"Force-updating artifacts for model {model_name} from: {full_gs_url}"
+        )
         download_public_file(full_gs_url, model_dir)
     elif shark_args.force_update_tank == True:
-        print(f"Force-updating artifacts for model {model_name}...")
+        print(
+            f"Force-updating artifacts for model {model_name} from: {full_gs_url}"
+        )
         download_public_file(full_gs_url, model_dir)
     else:
         if not _internet_connected():

shark/shark_importer.py

@@ -4,6 +4,17 @@
 import sys
 import tempfile
 import os
+import hashlib
+
+
+def create_hash(file_name):
+    with open(file_name, "rb") as f:
+        file_hash = hashlib.blake2b()
+        while chunk := f.read(2**20):
+            file_hash.update(chunk)
+
+    return file_hash.hexdigest()
+
 
 # List of the supported frontends.
 supported_frontends = {
@@ -140,6 +151,7 @@ class SharkImporter:
         outputs_name = "golden_out.npz"
         func_file_name = "function_name"
         model_name_mlir = model_name + "_" + self.frontend + ".mlir"
+        print(f"saving {model_name_mlir} to {dir}")
         try:
             inputs = [x.cpu().detach() for x in inputs]
         except AttributeError:
@@ -150,11 +162,11 @@ class SharkImporter:
         np.savez(os.path.join(dir, inputs_name), *inputs)
         np.savez(os.path.join(dir, outputs_name), *outputs)
         np.save(os.path.join(dir, func_file_name), np.array(func_name))
-
         if self.frontend == "torch":
             with open(os.path.join(dir, model_name_mlir), "wb") as mlir_file:
                 mlir_file.write(mlir_data)
-
+            mlir_hash = create_hash(os.path.join(dir, model_name_mlir))
+            np.save(os.path.join(dir, "hash"), np.array(mlir_hash))
         return
 
     def import_debug(
@@ -377,7 +389,10 @@ def import_with_fx(
 
     golden_values = None
     if debug:
-        golden_values = model(*inputs)
+        try:
+            golden_values = model(*inputs)
+        except:
+            golden_values = None
     # TODO: Control the decompositions.
     fx_g = make_fx(
         model,

tank/examples/MiniLM_tf/huggingface_MiniLM_run.py

@@ -63,7 +63,7 @@ if __name__ == "__main__":
     # Compile the model using IREE
     backend = "dylib-llvm-aot"
     args = [
-        "--iree-llvm-target-cpu-features=host",
+        "--iree-llvmcpu-target-cpu-features=host",
         "--iree-mhlo-demote-i64-to-i32=false",
         "--iree-flow-demote-i64-to-i32",
     ]

tank/examples/bert_fine_tuning/bert_fine_tune_tf.py

@@ -136,7 +136,7 @@ if __name__ == "__main__":
     backend = "dylib-llvm-aot"
     if backend == "dylib-llvm-aot":
         args = [
-            "--iree-llvm-target-cpu-features=host",
+            "--iree-llvmcpu-target-cpu-features=host",
             "--iree-mhlo-demote-i64-to-i32=false",
             "--iree-flow-demote-i64-to-i32",
         ]

tank/examples/bert_tf/bert_large_run.py

@@ -83,7 +83,7 @@ if __name__ == "__main__":
     # Compile the model using IREE
     backend = "dylib-llvm-aot"
     args = [
-        "--iree-llvm-target-cpu-features=host",
+        "--iree-llvmcpu-target-cpu-features=host",
         "--iree-mhlo-demote-i64-to-i32=false",
         "--iree-stream-resource-index-bits=64",
         "--iree-vm-target-index-bits=64",

tank/examples/bert_tf/bert_small_run.py

@@ -79,7 +79,7 @@ if __name__ == "__main__":
     # Compile the model using IREE
     backend = "dylib-llvm-aot"
     args = [
-        "--iree-llvm-target-cpu-features=host",
+        "--iree-llvmcpu-target-cpu-features=host",
         "--iree-mhlo-demote-i64-to-i32=false",
         "--iree-flow-demote-i64-to-i32",
     ]

tank/torch_sd_list.csv

@@ -0,0 +1,4 @@
+model_name, use_tracing, model_type, dynamic, param_count, tags, notes
+stabilityai/stable-diffusion-2-1-base,True,stable_diffusion,False,??M,"stable diffusion 2.1 base, LLM, Text to image", N/A
+stabilityai/stable-diffusion-2-1,True,stable_diffusion,False,??M,"stable diffusion 2.1 base, LLM, Text to image", N/A
+prompthero/openjourney,True,stable_diffusion,False,??M,"stable diffusion 2.1 base, LLM, Text to image", N/A