[Web] remove txt2img ui dependencies from png import metadata (#1275 )

Add auto-tuner to SD apps (#1291 )
Add documentation for using SHARK with AI-Render (#1296 )
2026-04-20 03:00:34 -04:00 · 2023-04-12 07:32:47 -10:00 · 2023-04-12 09:21:17 -07:00 · 2023-04-12 03:09:34 -10:00 · 2023-04-11 17:59:09 -07:00 · 2023-04-11 15:34:25 -05:00
29 changed files with 748 additions and 370 deletions
--- a/apps/stable_diffusion/scripts/init.py
+++ b/apps/stable_diffusion/scripts/init.py
@@ -1,4 +1,3 @@
-from apps.stable_diffusion.scripts.img2img import img2img_inf
 from apps.stable_diffusion.scripts.inpaint import inpaint_inf
 from apps.stable_diffusion.scripts.outpaint import outpaint_inf
 from apps.stable_diffusion.scripts.upscaler import upscaler_inf
--- a/apps/stable_diffusion/scripts/img2img.py
+++ b/apps/stable_diffusion/scripts/img2img.py
@@ -7,6 +7,7 @@ from apps.stable_diffusion.src import (
    args,
    Image2ImagePipeline,
    StencilPipeline,
+    resize_stencil,
    get_schedulers,
    set_init_device_flags,
    utils,
@@ -16,273 +17,6 @@ from apps.stable_diffusion.src import (
 from apps.stable_diffusion.src.utils import get_generation_text_info


-# 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
-
-
-# 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 in the range of [128, 768] 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 < 128:
-        n_size = 128
-        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
-
-    min_size = min(width, height)
-    if min_size > 768:
-        n_size = 768
-        if width == min_size:
-            height = n_size
-            width = n_size * aspect_ratio
-        else:
-            width = n_size
-            height = 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,
-    height: int,
-    width: int,
-    steps: int,
-    strength: float,
-    guidance_scale: float,
-    seed: int,
-    batch_count: int,
-    batch_size: int,
-    scheduler: str,
-    custom_model: str,
-    hf_model_id: str,
-    precision: str,
-    device: str,
-    max_length: int,
-    use_stencil: str,
-    save_metadata_to_json: bool,
-    save_metadata_to_png: bool,
-    lora_weights: str,
-    lora_hf_id: str,
-    ondemand: bool,
-):
-    from apps.stable_diffusion.web.ui.utils import (
-        get_custom_model_pathfile,
-        get_custom_vae_or_lora_weights,
-        Config,
-    )
-    import apps.stable_diffusion.web.utils.global_obj as global_obj
-    from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
-        SD_STATE_CANCEL,
-    )
-
-    args.prompts = [prompt]
-    args.negative_prompts = [negative_prompt]
-    args.guidance_scale = guidance_scale
-    args.seed = seed
-    args.steps = steps
-    args.strength = strength
-    args.scheduler = scheduler
-    args.img_path = "not none"
-    args.ondemand = ondemand
-
-    if init_image is None:
-        return None, "An Initial Image is required"
-    image = init_image.convert("RGB")
-
-    # set ckpt_loc and hf_model_id.
-    args.ckpt_loc = ""
-    args.hf_model_id = ""
-    if custom_model == "None":
-        if not hf_model_id:
-            return (
-                None,
-                "Please provide either custom model or huggingface model ID, both must not be empty",
-            )
-        args.hf_model_id = hf_model_id
-    elif ".ckpt" in custom_model or ".safetensors" in custom_model:
-        args.ckpt_loc = get_custom_model_pathfile(custom_model)
-    else:
-        args.hf_model_id = custom_model
-
-    args.use_lora = get_custom_vae_or_lora_weights(
-        lora_weights, lora_hf_id, "lora"
-    )
-
-    args.save_metadata_to_json = save_metadata_to_json
-    args.write_metadata_to_png = save_metadata_to_png
-
-    use_stencil = None if use_stencil == "None" else use_stencil
-    args.use_stencil = use_stencil
-    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(
-                f"SharkEulerDiscrete scheduler not supported. Switching to PNDM scheduler"
-            )
-            args.scheduler = "PNDM"
-        else:
-            sys.exit(
-                "Img2Img works best with PNDM scheduler. Other schedulers are not supported yet."
-            )
-    cpu_scheduling = not args.scheduler.startswith("Shark")
-    args.precision = precision
-    dtype = torch.float32 if precision == "fp32" else torch.half
-    new_config_obj = Config(
-        "img2img",
-        args.hf_model_id,
-        args.ckpt_loc,
-        precision,
-        batch_size,
-        max_length,
-        height,
-        width,
-        device,
-        use_lora=args.use_lora,
-        use_stencil=use_stencil,
-    )
-    if (
-        not global_obj.get_sd_obj()
-        or global_obj.get_cfg_obj() != new_config_obj
-    ):
-        global_obj.clear_cache()
-        global_obj.set_cfg_obj(new_config_obj)
-        args.batch_count = batch_count
-        args.batch_size = batch_size
-        args.max_length = max_length
-        args.height = height
-        args.width = width
-        args.device = device.split("=>", 1)[1].strip()
-        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
-            if args.hf_model_id
-            else "stabilityai/stable-diffusion-2-1-base"
-        )
-        global_obj.set_schedulers(get_schedulers(model_id))
-        scheduler_obj = global_obj.get_scheduler(args.scheduler)
-
-        if use_stencil is not None:
-            args.use_tuned = False
-            global_obj.set_sd_obj(
-                StencilPipeline.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,
-                    low_cpu_mem_usage=args.low_cpu_mem_usage,
-                    use_stencil=use_stencil,
-                    debug=args.import_debug if args.import_mlir else False,
-                    use_lora=args.use_lora,
-                    ondemand=args.ondemand,
-                )
-            )
-        else:
-            global_obj.set_sd_obj(
-                Image2ImagePipeline.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,
-                    low_cpu_mem_usage=args.low_cpu_mem_usage,
-                    debug=args.import_debug if args.import_mlir else False,
-                    use_lora=args.use_lora,
-                    ondemand=args.ondemand,
-                )
-            )
-
-    global_obj.set_sd_scheduler(args.scheduler)
-
-    start_time = time.time()
-    global_obj.get_sd_obj().log = ""
-    generated_imgs = []
-    seeds = []
-    img_seed = utils.sanitize_seed(seed)
-    extra_info = {"STRENGTH": strength}
-    text_output = ""
-    for current_batch in range(batch_count):
-        if current_batch > 0:
-            img_seed = utils.sanitize_seed(-1)
-        out_imgs = global_obj.get_sd_obj().generate_images(
-            prompt,
-            negative_prompt,
-            image,
-            batch_size,
-            height,
-            width,
-            steps,
-            strength,
-            guidance_scale,
-            img_seed,
-            args.max_length,
-            dtype,
-            args.use_base_vae,
-            cpu_scheduling,
-            use_stencil=use_stencil,
-        )
-        seeds.append(img_seed)
-        total_time = time.time() - start_time
-        text_output = get_generation_text_info(seeds, device)
-        text_output += "\n" + global_obj.get_sd_obj().log
-        text_output += f"\nTotal image(s) generation time: {total_time:.4f}sec"
-
-        if global_obj.get_sd_status() == SD_STATE_CANCEL:
-            break
-        else:
-            save_output_img(out_imgs[0], img_seed, extra_info)
-            generated_imgs.extend(out_imgs)
-            yield generated_imgs, text_output
-
-    return generated_imgs, text_output
-
-
 def main():
    if args.clear_all:
        clear_all()
--- a/apps/stable_diffusion/scripts/inpaint.py
+++ b/apps/stable_diffusion/scripts/inpaint.py
@@ -64,6 +64,9 @@ def inpaint_inf(
    args.img_path = "not none"
    args.mask_path = "not none"
    args.ondemand = ondemand
+    if ondemand and batch_count > 1:
+        print("Low VRAM mode currently only supports 1 batch count.")
+        batch_count = 1

    # set ckpt_loc and hf_model_id.
    args.ckpt_loc = ""
@@ -101,9 +104,11 @@ def inpaint_inf(
        device,
        use_lora=args.use_lora,
        use_stencil=None,
+        ondemand=ondemand,
    )
    if (
-        not global_obj.get_sd_obj()
+        args.ondemand
+        or not global_obj.get_sd_obj()
        or global_obj.get_cfg_obj() != new_config_obj
    ):
        global_obj.clear_cache()
--- a/apps/stable_diffusion/scripts/outpaint.py
+++ b/apps/stable_diffusion/scripts/outpaint.py
@@ -65,6 +65,9 @@ def outpaint_inf(
    args.steps = steps
    args.scheduler = scheduler
    args.img_path = "not none"
+    if ondemand:
+        print("Outpainting is not supporting ondemand yet.")
+        ondemand = False
    args.ondemand = ondemand

    # set ckpt_loc and hf_model_id.
--- a/apps/stable_diffusion/scripts/tuner.py
+++ b/apps/stable_diffusion/scripts/tuner.py
@@ -0,0 +1,114 @@
+import os
+from pathlib import Path
+from shark_tuner.codegen_tuner import SharkCodegenTuner
+from shark_tuner.iree_utils import (
+    dump_dispatches,
+    create_context,
+    export_module_to_mlir_file,
+)
+from shark_tuner.model_annotation import model_annotation
+from apps.stable_diffusion.src.utils.stable_args import args
+from apps.stable_diffusion.src.utils.utils import set_init_device_flags
+from apps.stable_diffusion.src.utils.sd_annotation import (
+    get_device_args,
+    load_winograd_configs,
+)
+from apps.stable_diffusion.src.models import SharkifyStableDiffusionModel
+
+
+def load_mlir_module():
+    sd_model = SharkifyStableDiffusionModel(
+        args.hf_model_id,
+        args.ckpt_loc,
+        args.custom_vae,
+        args.precision,
+        max_len=args.max_length,
+        batch_size=args.batch_size,
+        height=args.height,
+        width=args.width,
+        use_base_vae=args.use_base_vae,
+        use_tuned=False,
+        low_cpu_mem_usage=args.low_cpu_mem_usage,
+        return_mlir=True,
+    )
+
+    if args.annotation_model == "unet":
+        mlir_module = sd_model.unet()
+        model_name = sd_model.model_name["unet"]
+    elif args.annotation_model == "vae":
+        mlir_module = sd_model.vae()
+        model_name = sd_model.model_name["vae"]
+    else:
+        raise ValueError(
+            f"{args.annotation_model} is not supported for tuning."
+        )
+
+    return mlir_module, model_name
+
+
+def main():
+    args.use_tuned = False
+    set_init_device_flags()
+    mlir_module, model_name = load_mlir_module()
+
+    # Get device and device specific arguments
+    device, device_spec_args = get_device_args()
+    device_spec = ""
+    if device_spec_args:
+        device_spec = device_spec_args[-1].split("=")[-1].strip()
+        if device == "vulkan":
+            device_spec = device_spec.split("-")[0]
+
+    # Add winograd annotation for vulkan device
+    winograd_config = (
+        load_winograd_configs()
+        if device == "vulkan" and args.annotation_model in ["unet", "vae"]
+        else ""
+    )
+    with create_context() as ctx:
+        input_module = model_annotation(
+            ctx,
+            input_contents=mlir_module,
+            config_path=winograd_config,
+            search_op="conv",
+            winograd=True,
+        )
+
+    # Dump model dispatches
+    if device == "vulkan" and device_spec == "rdna3":
+        device = "vulkan/RX 7900"
+    generates_dir = Path.home() / "tmp"
+    if not os.path.exists(generates_dir):
+        os.makedirs(generates_dir)
+    dump_mlir = generates_dir / "temp.mlir"
+    dispatch_dir = generates_dir / f"{model_name}_{device_spec}_dispatches"
+    export_module_to_mlir_file(input_module, dump_mlir)
+    dump_dispatches(dump_mlir, device, dispatch_dir, False)
+
+    # Tune each dispatch
+    dtype = "f16" if args.precision == "fp16" else "f32"
+    config_filename = f"{model_name}_{device_spec}_configs.json"
+
+    for f_path in os.listdir(dispatch_dir):
+        if not f_path.endswith(".mlir"):
+            continue
+
+        model_dir = os.path.join(dispatch_dir, f_path)
+
+        tuner = SharkCodegenTuner(
+            model_dir,
+            device,
+            "random",
+            args.num_iters,
+            args.tuned_config_dir,
+            dtype,
+            args.search_op,
+            batch_size=1,
+            config_filename=config_filename,
+            use_dispatch=True,
+        )
+        tuner.tune()
+
+
+if __name__ == "__main__":
+    main()
--- a/apps/stable_diffusion/scripts/upscaler.py
+++ b/apps/stable_diffusion/scripts/upscaler.py
@@ -58,6 +58,9 @@ def upscaler_inf(
    args.steps = steps
    args.scheduler = scheduler
    args.ondemand = ondemand
+    if ondemand and batch_count > 1:
+        print("Low VRAM mode currently only supports 1 batch count.")
+        batch_count = 1

    if init_image is None:
        return None, "An Initial Image is required"
@@ -101,9 +104,11 @@ def upscaler_inf(
        device,
        use_lora=args.use_lora,
        use_stencil=None,
+        ondemand=ondemand,
    )
    if (
-        not global_obj.get_sd_obj()
+        args.ondemand
+        or not global_obj.get_sd_obj()
        or global_obj.get_cfg_obj() != new_config_obj
    ):
        global_obj.clear_cache()
--- a/apps/stable_diffusion/src/init.py
+++ b/apps/stable_diffusion/src/init.py
@@ -5,6 +5,7 @@ from apps.stable_diffusion.src.utils import (
    get_available_devices,
    clear_all,
    save_output_img,
+    resize_stencil,
 )
 from apps.stable_diffusion.src.pipelines import (
    Text2ImagePipeline,
--- a/apps/stable_diffusion/src/models/model_wrappers.py
+++ b/apps/stable_diffusion/src/models/model_wrappers.py
@@ -82,6 +82,7 @@ class SharkifyStableDiffusionModel:
        use_stencil: str = None,
        use_lora: str = "",
        use_quantize: str = None,
+        return_mlir: bool = False,
    ):
        self.check_params(max_len, width, height)
        self.max_len = max_len
@@ -147,6 +148,7 @@ class SharkifyStableDiffusionModel:
        self.base_model_id = fetch_and_update_base_model_id(self.model_to_run)
        if self.base_model_id != "" and args.ckpt_loc != "":
            args.hf_model_id = self.base_model_id
+        self.return_mlir = return_mlir

    def get_extended_name_for_all_model(self):
        model_name = {}
@@ -211,8 +213,8 @@ class SharkifyStableDiffusionModel:

        vae_encode = VaeEncodeModel()
        inputs = tuple(self.inputs["vae_encode"])
-        is_f16 = True if self.precision == "fp16" else False
-        shark_vae_encode = compile_through_fx(
+        is_f16 = True if not self.is_upscaler and self.precision == "fp16" else False
+        shark_vae_encode, vae_encode_mlir = compile_through_fx(
            vae_encode,
            inputs,
            is_f16=is_f16,
@@ -221,7 +223,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("vae", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_vae_encode
+        return shark_vae_encode, vae_encode_mlir

    def get_vae(self):
        class VaeModel(torch.nn.Module):
@@ -265,7 +267,7 @@ class SharkifyStableDiffusionModel:
        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(
+        shark_vae, vae_mlir = compile_through_fx(
            vae,
            inputs,
            is_f16=is_f16,
@@ -277,7 +279,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("vae", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_vae
+        return shark_vae, vae_mlir

    def get_controlled_unet(self):
        class ControlledUnetModel(torch.nn.Module):
@@ -322,7 +324,7 @@ class SharkifyStableDiffusionModel:

        inputs = tuple(self.inputs["unet"])
        input_mask = [True, True, True, False, True, True, True, True, True, True, True, True, True, True, True, True, True,]
-        shark_controlled_unet = compile_through_fx(
+        shark_controlled_unet, controlled_unet_mlir = compile_through_fx(
            unet,
            inputs,
            model_name=self.model_name["stencil_unet"],
@@ -332,7 +334,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("unet", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_controlled_unet
+        return shark_controlled_unet, controlled_unet_mlir

    def get_control_net(self):
        class StencilControlNetModel(torch.nn.Module):
@@ -376,7 +378,7 @@ class SharkifyStableDiffusionModel:

        inputs = tuple(self.inputs["stencil_adaptor"])
        input_mask = [True, True, True, True]
-        shark_cnet = compile_through_fx(
+        shark_cnet, cnet_mlir = compile_through_fx(
            scnet,
            inputs,
            model_name=self.model_name["stencil_adaptor"],
@@ -386,7 +388,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("unet", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_cnet
+        return shark_cnet, cnet_mlir

    def get_unet(self):
        class UnetModel(torch.nn.Module):
@@ -432,7 +434,7 @@ class SharkifyStableDiffusionModel:
                save_dir,
                exist_ok=True,
            )
-        shark_unet = compile_through_fx(
+        shark_unet, unet_mlir = compile_through_fx(
            unet,
            inputs,
            model_name=self.model_name["unet"],
@@ -445,7 +447,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("unet", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_unet
+        return shark_unet, unet_mlir

    def get_unet_upscaler(self):
        class UnetModel(torch.nn.Module):
@@ -473,7 +475,7 @@ class SharkifyStableDiffusionModel:
        is_f16 = True if self.precision == "fp16" else False
        inputs = tuple(self.inputs["unet"])
        input_mask = [True, True, True, False]
-        shark_unet = compile_through_fx(
+        shark_unet, unet_mlir = compile_through_fx(
            unet,
            inputs,
            model_name=self.model_name["unet"],
@@ -483,7 +485,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("unet", precision=self.precision),
            base_model_id=self.base_model_id,
        )
-        return shark_unet
+        return shark_unet, unet_mlir

    def get_clip(self):
        class CLIPText(torch.nn.Module):
@@ -507,7 +509,7 @@ class SharkifyStableDiffusionModel:
                save_dir,
                exist_ok=True,
            )
-        shark_clip = compile_through_fx(
+        shark_clip, clip_mlir = compile_through_fx(
            clip_model,
            tuple(self.inputs["clip"]),
            model_name=self.model_name["clip"],
@@ -517,7 +519,7 @@ class SharkifyStableDiffusionModel:
            extra_args=get_opt_flags("clip", precision="fp32"),
            base_model_id=self.base_model_id,
        )
-        return shark_clip
+        return shark_clip, clip_mlir

    def process_custom_vae(self):
        custom_vae = self.custom_vae.lower()
@@ -555,28 +557,32 @@ class SharkifyStableDiffusionModel:
    def vae_encode(self):
        # Fetch vmfb for the model if present
        vmfb = fetch_vmfb("vae_encode", self.model_name["vae_encode"], self.precision)
-        if vmfb:
+        if vmfb and not self.return_mlir:
            return vmfb

        try:
            self.inputs["vae_encode"] = self.get_input_info_for(base_models["vae_encode"])
-            compiled_vae_encode = self.get_vae_encode()
+            compiled_vae_encode, vae_encode_mlir = self.get_vae_encode()

            check_compilation(compiled_vae_encode, "Vae Encode")
+            if self.return_mlir:
+                return vae_encode_mlir
            return compiled_vae_encode
        except Exception as e:
            sys.exit(e)

    def clip(self):
        vmfb = fetch_vmfb("clip", self.model_name["clip"], self.precision)
-        if vmfb:
+        if vmfb and not self.return_mlir:
            return vmfb

        try:
            self.inputs["clip"] = self.get_input_info_for(base_models["clip"])
-            compiled_clip = self.get_clip()
+            compiled_clip, clip_mlir = self.get_clip()

            check_compilation(compiled_clip, "Clip")
+            if self.return_mlir:
+                return clip_mlir
            return compiled_clip
        except Exception as e:
            sys.exit(e)
@@ -584,7 +590,7 @@ class SharkifyStableDiffusionModel:
    def unet(self):
        model = "stencil_unet" if self.use_stencil is not None else "unet"
        vmfb = fetch_vmfb(model, self.model_name[model], self.precision)
-        if vmfb:
+        if vmfb and not self.return_mlir:
            return vmfb

        try:
@@ -593,14 +599,14 @@ class SharkifyStableDiffusionModel:

            if self.base_model_id != "":
                self.inputs["unet"] = self.get_input_info_for(unet_inputs[self.base_model_id])
-                compiled_unet = self.compile_unet_variants(model)
+                compiled_unet, unet_mlir = self.compile_unet_variants(model)
            else:
                for model_id in unet_inputs:
                    self.base_model_id = model_id
                    self.inputs["unet"] = self.get_input_info_for(unet_inputs[model_id])

                    try:
-                        compiled_unet = self.compile_unet_variants(model)
+                        compiled_unet, unet_mlir = self.compile_unet_variants(model)
                    except Exception as e:
                        print(e)
                        print("Retrying with a different base model configuration")
@@ -618,13 +624,15 @@ class SharkifyStableDiffusionModel:
                    break

            check_compilation(compiled_unet, "Unet")
+            if self.return_mlir:
+                return unet_mlir
            return compiled_unet
        except Exception as e:
            sys.exit(e)

    def vae(self):
        vmfb = fetch_vmfb("vae", self.model_name["vae"], self.precision)
-        if vmfb:
+        if vmfb and not self.return_mlir:
            return vmfb

        try:
@@ -634,24 +642,28 @@ class SharkifyStableDiffusionModel:
            is_base_vae = self.base_vae
            if self.is_upscaler:
                self.base_vae = True
-            compiled_vae = self.get_vae()
+            compiled_vae, vae_mlir = self.get_vae()
            self.base_vae = is_base_vae

            check_compilation(compiled_vae, "Vae")
+            if self.return_mlir:
+                return vae_mlir
            return compiled_vae
        except Exception as e:
            sys.exit(e)

    def controlnet(self):
        vmfb = fetch_vmfb("stencil_adaptor", self.model_name["stencil_adaptor"], self.precision)
-        if vmfb:
+        if vmfb and not self.return_mlir:
            return vmfb

        try:
            self.inputs["stencil_adaptor"] = self.get_input_info_for(base_models["stencil_adaptor"])
-            compiled_stencil_adaptor = self.get_control_net()
+            compiled_stencil_adaptor, controlnet_mlir = self.get_control_net()

            check_compilation(compiled_stencil_adaptor, "Stencil")
+            if self.return_mlir:
+                return controlnet_mlir
            return compiled_stencil_adaptor
        except Exception as e:
            sys.exit(e)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py
@@ -155,7 +155,9 @@ class Image2ImagePipeline(StableDiffusionPipeline):
        generator = torch.manual_seed(seed)

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # guidance scale as a float32 tensor.
        guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py
@@ -407,7 +407,9 @@ class InpaintPipeline(StableDiffusionPipeline):
        )

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # guidance scale as a float32 tensor.
        guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py
@@ -408,7 +408,9 @@ class OutpaintPipeline(StableDiffusionPipeline):
        )

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # guidance scale as a float32 tensor.
        guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
@@ -537,8 +539,6 @@ class OutpaintPipeline(StableDiffusionPipeline):
                    cpu_scheduling=cpu_scheduling,
                )
                all_imgs.extend(imgs)
-            if self.ondemand:
-                self.unload_vae()

            res_img = all_imgs[0].resize(
                (image_to_process.width, image_to_process.height)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py
@@ -229,7 +229,9 @@ class StencilPipeline(StableDiffusionPipeline):
        generator = torch.manual_seed(seed)

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # guidance scale as a float32 tensor.
        guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py
@@ -111,7 +111,9 @@ class Text2ImagePipeline(StableDiffusionPipeline):
        )

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # guidance scale as a float32 tensor.
        guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_upscaler.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_upscaler.py
@@ -256,7 +256,9 @@ class UpscalerPipeline(StableDiffusionPipeline):
        generator = torch.manual_seed(seed)

        # Get text embeddings with weight emphasis from prompts
-        text_embeddings = self.encode_prompts_weight(prompts, neg_prompts)
+        text_embeddings = self.encode_prompts_weight(
+            prompts, neg_prompts, max_length
+        )

        # 4. Preprocess image
        image = preprocess(image).to(dtype)
--- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py
+++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py
@@ -57,7 +57,7 @@ class StableDiffusionPipeline:
        self.vae = None
        self.text_encoder = None
        self.unet = None
-        self.tokenizer = get_tokenizer()
+        self.model_max_length = 77
        self.scheduler = scheduler
        # TODO: Implement using logging python utility.
        self.log = ""
@@ -66,6 +66,13 @@ class StableDiffusionPipeline:
        self.import_mlir = import_mlir
        self.use_lora = use_lora
        self.ondemand = ondemand
+        # TODO: Find a better workaround for fetching base_model_id early enough for CLIPTokenizer.
+        try:
+            self.tokenizer = get_tokenizer()
+        except:
+            self.load_unet()
+            self.unload_unet()
+            self.tokenizer = get_tokenizer()

    def load_clip(self):
        if self.text_encoder is not None:
@@ -146,7 +153,8 @@ class StableDiffusionPipeline:
        clip_inf_start = time.time()
        text_embeddings = self.text_encoder("forward", (text_input,))
        clip_inf_time = (time.time() - clip_inf_start) * 1000
-        # self.unload_clip()
+        if self.ondemand:
+            self.unload_clip()
        self.log += f"\nClip Inference time (ms) = {clip_inf_time:.3f}"

        return text_embeddings
@@ -337,6 +345,7 @@ class StableDiffusionPipeline:
        self,
        prompt,
        negative_prompt,
+        model_max_length,
        do_classifier_free_guidance=True,
        max_embeddings_multiples=1,
        num_images_per_prompt=1,
@@ -349,6 +358,8 @@ class StableDiffusionPipeline:
            negative_prompt (`str` or `List[str]`):
                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
                if `guidance_scale` is less than `1`).
+            model_max_length (int):
+                SHARK: pass the max length instead of relying on pipe.tokenizer.model_max_length
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not,
                SHARK: must be set to True as we always expect neg embeddings (defaulted to True)
@@ -360,7 +371,8 @@ class StableDiffusionPipeline:
                SHARK: num_images_per_prompt is not used (defaulted to 1)
        """

-        # SHARK: Load the clip and prepare inference time
+        # SHARK: Save model_max_length, load the clip and init inference time
+        self.model_max_length = model_max_length
        self.load_clip()
        clip_inf_start = time.time()

@@ -399,7 +411,8 @@ class StableDiffusionPipeline:

        # SHARK: Report clip inference time
        clip_inf_time = (time.time() - clip_inf_start) * 1000
-        # self.unload_clip()
+        if self.ondemand:
+            self.unload_clip()
        self.log += f"\nClip Inference time (ms) = {clip_inf_time:.3f}"

        return text_embeddings.numpy()
@@ -689,9 +702,7 @@ def get_weighted_text_embeddings(
        skip_weighting (`bool`, *optional*, defaults to `False`):
            Skip the weighting. When the parsing is skipped, it is forced True.
    """
-    max_length = (
-        pipe.tokenizer.model_max_length - 2
-    ) * max_embeddings_multiples + 2
+    max_length = (pipe.model_max_length - 2) * max_embeddings_multiples + 2
    if isinstance(prompt, str):
        prompt = [prompt]

@@ -733,12 +744,10 @@ def get_weighted_text_embeddings(

    max_embeddings_multiples = min(
        max_embeddings_multiples,
-        (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1,
+        (max_length - 1) // (pipe.model_max_length - 2) + 1,
    )
    max_embeddings_multiples = max(1, max_embeddings_multiples)
-    max_length = (
-        pipe.tokenizer.model_max_length - 2
-    ) * max_embeddings_multiples + 2
+    max_length = (pipe.model_max_length - 2) * max_embeddings_multiples + 2

    # pad the length of tokens and weights
    bos = pipe.tokenizer.bos_token_id
@@ -750,7 +759,7 @@ def get_weighted_text_embeddings(
        bos,
        eos,
        no_boseos_middle=no_boseos_middle,
-        chunk_length=pipe.tokenizer.model_max_length,
+        chunk_length=pipe.model_max_length,
    )
    # prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=pipe.device)
    prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device="cpu")
@@ -762,7 +771,7 @@ def get_weighted_text_embeddings(
            bos,
            eos,
            no_boseos_middle=no_boseos_middle,
-            chunk_length=pipe.tokenizer.model_max_length,
+            chunk_length=pipe.model_max_length,
        )
        # uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=pipe.device)
        uncond_tokens = torch.tensor(
@@ -773,7 +782,7 @@ def get_weighted_text_embeddings(
    text_embeddings = get_unweighted_text_embeddings(
        pipe,
        prompt_tokens,
-        pipe.tokenizer.model_max_length,
+        pipe.model_max_length,
        no_boseos_middle=no_boseos_middle,
    )
    # prompt_weights = torch.tensor(prompt_weights, dtype=text_embeddings.dtype, device=pipe.device)
@@ -784,7 +793,7 @@ def get_weighted_text_embeddings(
        uncond_embeddings = get_unweighted_text_embeddings(
            pipe,
            uncond_tokens,
-            pipe.tokenizer.model_max_length,
+            pipe.model_max_length,
            no_boseos_middle=no_boseos_middle,
        )
        # uncond_weights = torch.tensor(uncond_weights, dtype=uncond_embeddings.dtype, device=pipe.device)
--- a/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py
+++ b/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py
@@ -89,7 +89,7 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):

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

            step_model = SchedulerStepModel()
-            self.step_model = compile_through_fx(
+            self.step_model, _ = compile_through_fx(
                step_model,
                (example_output, example_sigma, example_latent, example_dt),
                model_name=f"euler_step_{BATCH_SIZE}_{args.height}_{args.width}"
--- a/apps/stable_diffusion/src/utils/init.py
+++ b/apps/stable_diffusion/src/utils/init.py
@@ -34,4 +34,5 @@ from apps.stable_diffusion.src.utils.utils import (
    save_output_img,
    get_generation_text_info,
    update_lora_weight,
+    resize_stencil,
 )
--- a/apps/stable_diffusion/src/utils/stable_args.py
+++ b/apps/stable_diffusion/src/utils/stable_args.py
@@ -539,6 +539,31 @@ p.add_argument(
    action=argparse.BooleanOptionalAction,
    help="Save annotated mlir file",
 )
+##############################################################################
+### SD model auto-tuner flags
+##############################################################################
+
+p.add_argument(
+    "--tuned_config_dir",
+    type=path_expand,
+    default="./",
+    help="Directory to save the tuned config file",
+)
+
+p.add_argument(
+    "--num_iters",
+    type=int,
+    default=400,
+    help="Number of iterations for tuning",
+)
+
+p.add_argument(
+    "--search_op",
+    type=str,
+    default="all",
+    help="Op to be optimized, options are matmul, bmm, conv and all",
+)
+

 args, unknown = p.parse_known_args()
 if args.import_debug:
--- a/apps/stable_diffusion/src/utils/utils.py
+++ b/apps/stable_diffusion/src/utils/utils.py
@@ -3,6 +3,7 @@ import gc
 import json
 import re
 from PIL import PngImagePlugin
+from PIL import Image
 from datetime import datetime as dt
 from csv import DictWriter
 from pathlib import Path
@@ -135,9 +136,10 @@ def compile_through_fx(
            device=args.device,
            mlir_dialect="tm_tensor",
        )
-        del mlir_module
-        gc.collect()
-        return _compile_module(shark_module, model_name, extra_args)
+        return (
+            _compile_module(shark_module, model_name, extra_args),
+            mlir_module,
+        )

    del mlir_module
    gc.collect()
@@ -749,3 +751,46 @@ def get_generation_text_info(seeds, device):
    text_output += f"\nsize={args.height}x{args.width}, batch_count={args.batch_count}, batch_size={args.batch_size}, max_length={args.max_length}"

    return text_output
+
+
+# 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 in the range of [128, 768] 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 < 128:
+        n_size = 128
+        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
+
+    min_size = min(width, height)
+    if min_size > 768:
+        n_size = 768
+        if width == min_size:
+            height = n_size
+            width = n_size * aspect_ratio
+        else:
+            width = n_size
+            height = 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
--- a/apps/stable_diffusion/web/index.py
+++ b/apps/stable_diffusion/web/index.py
@@ -2,6 +2,7 @@ import os
 import sys
 import transformers
 from apps.stable_diffusion.src import args, clear_all
+import apps.stable_diffusion.web.utils.global_obj as global_obj

 if sys.platform == "darwin":
    os.environ["DYLD_LIBRARY_PATH"] = "/usr/local/lib"
@@ -11,18 +12,21 @@ if args.clear_all:

 if __name__ == "__main__":
    if args.api:
-        from apps.stable_diffusion.web.ui import txt2img_inf
+        from apps.stable_diffusion.web.ui import txt2img_inf, img2img_api
        from fastapi import FastAPI, APIRouter
        import uvicorn

+        # init global sd pipeline and config
+        global_obj._init()
+
        app = FastAPI()
-        app.add_api_route("/sdapi/txt2img", txt2img_inf, methods=["post"])
+        app.add_api_route("/sdapi/v1/txt2img", txt2img_inf, methods=["post"])
+        app.add_api_route("/sdapi/v1/img2img", img2img_api, methods=["post"])
        app.include_router(APIRouter())
-        uvicorn.run(app, host="0.0.0.0", port=args.server_port)
+        uvicorn.run(app, host="127.0.0.1", port=args.server_port)
        sys.exit(0)

    import gradio as gr
-    import apps.stable_diffusion.web.utils.global_obj as global_obj
    from apps.stable_diffusion.web.utils.gradio_configs import (
        clear_gradio_tmp_imgs_folder,
    )
--- a/apps/stable_diffusion/web/ui/init.py
+++ b/apps/stable_diffusion/web/ui/init.py
@@ -8,6 +8,8 @@ from apps.stable_diffusion.web.ui.txt2img_ui import (
    txt2img_sendto_upscaler,
 )
 from apps.stable_diffusion.web.ui.img2img_ui import (
+    img2img_api,
+    img2img_inf,
    img2img_web,
    img2img_gallery,
    img2img_init_image,
--- a/apps/stable_diffusion/web/ui/img2img_ui.py
+++ b/apps/stable_diffusion/web/ui/img2img_ui.py
@@ -1,8 +1,13 @@
 from pathlib import Path
 import os
+import torch
+import time
+import sys
 import gradio as gr
 from PIL import Image
-from apps.stable_diffusion.scripts import img2img_inf
+import base64
+from io import BytesIO
+from fastapi.exceptions import HTTPException
 from apps.stable_diffusion.src import args
 from apps.stable_diffusion.web.ui.utils import (
    available_devices,
@@ -13,6 +18,314 @@ from apps.stable_diffusion.web.ui.utils import (
    predefined_models,
    cancel_sd,
 )
+from apps.stable_diffusion.src import (
+    args,
+    Image2ImagePipeline,
+    StencilPipeline,
+    resize_stencil,
+    get_schedulers,
+    set_init_device_flags,
+    utils,
+    clear_all,
+    save_output_img,
+)
+from apps.stable_diffusion.src.utils import get_generation_text_info
+
+
+# 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 img2img_inf(
+    prompt: str,
+    negative_prompt: str,
+    init_image,
+    height: int,
+    width: int,
+    steps: int,
+    strength: float,
+    guidance_scale: float,
+    seed: int,
+    batch_count: int,
+    batch_size: int,
+    scheduler: str,
+    custom_model: str,
+    hf_model_id: str,
+    precision: str,
+    device: str,
+    max_length: int,
+    use_stencil: str,
+    save_metadata_to_json: bool,
+    save_metadata_to_png: bool,
+    lora_weights: str,
+    lora_hf_id: str,
+    ondemand: bool,
+):
+    from apps.stable_diffusion.web.ui.utils import (
+        get_custom_model_pathfile,
+        get_custom_vae_or_lora_weights,
+        Config,
+    )
+    import apps.stable_diffusion.web.utils.global_obj as global_obj
+    from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
+        SD_STATE_CANCEL,
+    )
+
+    args.prompts = [prompt]
+    args.negative_prompts = [negative_prompt]
+    args.guidance_scale = guidance_scale
+    args.seed = seed
+    args.steps = steps
+    args.strength = strength
+    args.scheduler = scheduler
+    args.img_path = "not none"
+    args.ondemand = ondemand
+    if ondemand and batch_count > 1:
+        print("Low VRAM mode currently only supports 1 batch count.")
+        batch_count = 1
+
+    if init_image is None:
+        return None, "An Initial Image is required"
+    image = init_image.convert("RGB")
+
+    # set ckpt_loc and hf_model_id.
+    args.ckpt_loc = ""
+    args.hf_model_id = ""
+    if custom_model == "None":
+        if not hf_model_id:
+            return (
+                None,
+                "Please provide either custom model or huggingface model ID, both must not be empty",
+            )
+        args.hf_model_id = hf_model_id
+    elif ".ckpt" in custom_model or ".safetensors" in custom_model:
+        args.ckpt_loc = get_custom_model_pathfile(custom_model)
+    else:
+        args.hf_model_id = custom_model
+
+    args.use_lora = get_custom_vae_or_lora_weights(
+        lora_weights, lora_hf_id, "lora"
+    )
+
+    args.save_metadata_to_json = save_metadata_to_json
+    args.write_metadata_to_png = save_metadata_to_png
+
+    use_stencil = None if use_stencil == "None" else use_stencil
+    args.use_stencil = use_stencil
+    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(
+                f"SharkEulerDiscrete scheduler not supported. Switching to PNDM scheduler"
+            )
+            args.scheduler = "PNDM"
+        else:
+            sys.exit(
+                "Img2Img works best with PNDM scheduler. Other schedulers are not supported yet."
+            )
+    cpu_scheduling = not args.scheduler.startswith("Shark")
+    args.precision = precision
+    dtype = torch.float32 if precision == "fp32" else torch.half
+    new_config_obj = Config(
+        "img2img",
+        args.hf_model_id,
+        args.ckpt_loc,
+        precision,
+        batch_size,
+        max_length,
+        height,
+        width,
+        device,
+        use_lora=args.use_lora,
+        use_stencil=use_stencil,
+        ondemand=ondemand,
+    )
+    if (
+        args.ondemand
+        or not global_obj.get_sd_obj()
+        or global_obj.get_cfg_obj() != new_config_obj
+    ):
+        global_obj.clear_cache()
+        global_obj.set_cfg_obj(new_config_obj)
+        args.batch_count = batch_count
+        args.batch_size = batch_size
+        args.max_length = max_length
+        args.height = height
+        args.width = width
+        args.device = device.split("=>", 1)[1].strip()
+        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
+            if args.hf_model_id
+            else "stabilityai/stable-diffusion-2-1-base"
+        )
+        global_obj.set_schedulers(get_schedulers(model_id))
+        scheduler_obj = global_obj.get_scheduler(args.scheduler)
+
+        if use_stencil is not None:
+            args.use_tuned = False
+            global_obj.set_sd_obj(
+                StencilPipeline.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,
+                    low_cpu_mem_usage=args.low_cpu_mem_usage,
+                    use_stencil=use_stencil,
+                    debug=args.import_debug if args.import_mlir else False,
+                    use_lora=args.use_lora,
+                    ondemand=args.ondemand,
+                )
+            )
+        else:
+            global_obj.set_sd_obj(
+                Image2ImagePipeline.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,
+                    low_cpu_mem_usage=args.low_cpu_mem_usage,
+                    debug=args.import_debug if args.import_mlir else False,
+                    use_lora=args.use_lora,
+                    ondemand=args.ondemand,
+                )
+            )
+
+    global_obj.set_sd_scheduler(args.scheduler)
+
+    start_time = time.time()
+    global_obj.get_sd_obj().log = ""
+    generated_imgs = []
+    seeds = []
+    img_seed = utils.sanitize_seed(seed)
+    extra_info = {"STRENGTH": strength}
+    text_output = ""
+    for current_batch in range(batch_count):
+        if current_batch > 0:
+            img_seed = utils.sanitize_seed(-1)
+        out_imgs = global_obj.get_sd_obj().generate_images(
+            prompt,
+            negative_prompt,
+            image,
+            batch_size,
+            height,
+            width,
+            steps,
+            strength,
+            guidance_scale,
+            img_seed,
+            args.max_length,
+            dtype,
+            args.use_base_vae,
+            cpu_scheduling,
+            use_stencil=use_stencil,
+        )
+        seeds.append(img_seed)
+        total_time = time.time() - start_time
+        text_output = get_generation_text_info(seeds, device)
+        text_output += "\n" + global_obj.get_sd_obj().log
+        text_output += f"\nTotal image(s) generation time: {total_time:.4f}sec"
+
+        if global_obj.get_sd_status() == SD_STATE_CANCEL:
+            break
+        else:
+            save_output_img(out_imgs[0], img_seed, extra_info)
+            generated_imgs.extend(out_imgs)
+            #  yield generated_imgs, text_output
+
+    return generated_imgs, text_output
+
+
+def decode_base64_to_image(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";", 1)[1].split(",", 1)[1]
+    try:
+        image = Image.open(BytesIO(base64.b64decode(encoding)))
+        return image
+    except Exception as err:
+        print(err)
+        raise HTTPException(status_code=500, detail="Invalid encoded image")
+
+
+def encode_pil_to_base64(images):
+    encoded_imgs = []
+    for image in images:
+        with BytesIO() as output_bytes:
+            if args.output_img_format.lower() == "png":
+                image.save(output_bytes, format="PNG")
+
+            elif args.output_img_format.lower() in ("jpg", "jpeg"):
+                image.save(output_bytes, format="JPEG")
+            else:
+                raise HTTPException(
+                    status_code=500, detail="Invalid image format"
+                )
+            bytes_data = output_bytes.getvalue()
+            encoded_imgs.append(base64.b64encode(bytes_data))
+    return encoded_imgs
+
+
+# Img2Img Rest API.
+def img2img_api(
+    InputData: dict,
+):
+    print(InputData)
+    init_image = decode_base64_to_image(InputData["init_images"][0])
+    res = img2img_inf(
+        InputData["prompt"],
+        InputData["negative_prompt"],
+        init_image,
+        InputData["height"],
+        InputData["width"],
+        InputData["steps"],
+        InputData["denoising_strength"],
+        InputData["cfg_scale"],
+        InputData["seed"],
+        batch_count=1,
+        batch_size=1,
+        scheduler="PNDM",
+        custom_model="None",
+        hf_model_id="stabilityai/stable-diffusion-2-1-base",
+        precision="fp16",
+        device=available_devices[0],
+        max_length=64,
+        use_stencil="None",
+        save_metadata_to_json=False,
+        save_metadata_to_png=False,
+        lora_weights="None",
+        lora_hf_id="",
+        ondemand=False,
+    )
+    return {
+        "images": encode_pil_to_base64(res[0]),
+        "parameters": {},
+        "info": res[1],
+    }


 with gr.Blocks(title="Image-to-Image") as img2img_web:
--- a/apps/stable_diffusion/web/ui/txt2img_ui.py
+++ b/apps/stable_diffusion/web/ui/txt2img_ui.py
@@ -13,6 +13,7 @@ from apps.stable_diffusion.web.ui.utils import (
    predefined_models,
    cancel_sd,
 )
+from apps.stable_diffusion.web.utils.png_metadata import import_png_metadata
 from apps.stable_diffusion.src import (
    args,
    Text2ImagePipeline,
@@ -68,6 +69,9 @@ def txt2img_inf(
    args.steps = steps
    args.scheduler = scheduler
    args.ondemand = ondemand
+    if ondemand and batch_count > 1:
+        print("Low VRAM mode currently only supports 1 batch count.")
+        batch_count = 1

    # set ckpt_loc and hf_model_id.
    args.ckpt_loc = ""
@@ -105,9 +109,11 @@ def txt2img_inf(
        device,
        use_lora=args.use_lora,
        use_stencil=None,
+        ondemand=ondemand,
    )
    if (
-        not global_obj.get_sd_obj()
+        args.ondemand
+        or not global_obj.get_sd_obj()
        or global_obj.get_cfg_obj() != new_config_obj
    ):
        global_obj.clear_cache()
@@ -447,14 +453,20 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
            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,
+                prompt,
+                negative_prompt,
+                steps,
+                scheduler,
+                guidance_scale,
+                seed,
+                width,
+                height,
+                custom_model,
+                hf_model_id,
            ],
            outputs=[
                png_info_img,
--- a/apps/stable_diffusion/web/ui/utils.py
+++ b/apps/stable_diffusion/web/ui/utils.py
@@ -24,6 +24,7 @@ class Config:
    device: str
    use_lora: str
    use_stencil: str
+    ondemand: str


 custom_model_filetypes = (
--- a/apps/stable_diffusion/web/utils/global_obj.py
+++ b/apps/stable_diffusion/web/utils/global_obj.py
@@ -43,18 +43,22 @@ def set_schedulers(value):


 def get_sd_obj():
+    global _sd_obj
    return _sd_obj


 def get_sd_status():
+    global _sd_obj
    return _sd_obj.status


 def get_cfg_obj():
+    global _config_obj
    return _config_obj


 def get_scheduler(key):
+    global _schedulers
    return _schedulers[key]


--- a/apps/stable_diffusion/web/utils/png_metadata.py
+++ b/apps/stable_diffusion/web/utils/png_metadata.py
@@ -1,18 +1,5 @@
 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,
@@ -75,7 +62,19 @@ def parse_generation_parameters(x: str):
    return res


-def import_png_metadata(pil_data):
+def import_png_metadata(
+    pil_data,
+    prompt,
+    negative_prompt,
+    steps,
+    sampler,
+    cfg_scale,
+    seed,
+    width,
+    height,
+    custom_model,
+    hf_model_id,
+):
    try:
        png_info = pil_data.info["parameters"]
        metadata = parse_generation_parameters(png_info)
@@ -110,39 +109,44 @@ def import_png_metadata(pil_data):
                    % 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"]),
-        }
+        negative_prompt = metadata["Negative prompt"]
+        steps = int(metadata["Steps"])
+        cfg_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] = ""
+            custom_model = png_custom_model
+            hf_model_id = ""
        if "Model" in metadata and png_hf_model_id:
-            outputs[custom_model] = "None"
-            outputs[hf_model_id] = png_hf_model_id
+            custom_model = "None"
+            hf_model_id = png_hf_model_id
        if "Prompt" in metadata:
-            outputs[prompt] = metadata["Prompt"]
+            prompt = metadata["Prompt"]
        if "Sampler" in metadata:
            if metadata["Sampler"] in scheduler_list_txt2img:
-                outputs[scheduler] = metadata["Sampler"]
+                sampler = 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,
-    }
+    return (
+        None,
+        prompt,
+        negative_prompt,
+        steps,
+        sampler,
+        cfg_scale,
+        seed,
+        width,
+        height,
+        custom_model,
+        hf_model_id,
+    )
--- a/docs/shark_sd_blender.md
+++ b/docs/shark_sd_blender.md
@@ -0,0 +1,75 @@
+# Overview
+
+This document is intended to provide a starting point for using SHARK stable diffusion with Blender. 
+
+We currently make use of the [AI-Render Plugin](https://github.com/benrugg/AI-Render) to integrate with Blender.
+
+## Setup SHARK and prerequisites:
+
+ * Download the latest SHARK SD webui .exe from [here](https://github.com/nod-ai/SHARK/releases) or follow instructions on the [README](https://github.com/nod-ai/SHARK#readme)
+ * Once you have the .exe where you would like SHARK to install, run the .exe from terminal/PowerShell with the `--api` flag:
+```
+## Run the .exe in API mode:
+.\shark_sd_<date>_<ver>.exe --api
+
+## For example:
+.\shark_sd_20230411_671.exe --api --server_port=8082
+
+## From a the base directory of a source clone of SHARK:
+./setup_venv.ps1
+python apps\stable_diffusion\web\index.py --api
+
+```
+
+Your local SD server should start and look something like this:
+![image](https://user-images.githubusercontent.com/87458719/231369758-e2c3c45a-eccc-4fe5-a788-4a3bf1ace1d1.png)
+
+ * Note: When running in api mode with `--api`, the .exe will not function as a webUI. Thus, the address in the terminal output will only be useful for API requests.
+
+### Install AI Render
+
+- Get AI Render on [Blender Market](https://blendermarket.com/products/ai-render) or [Gumroad](https://airender.gumroad.com/l/ai-render)
+- Open Blender, then go to Edit > Preferences > Add-ons > Install and then find the zip file
+- We will be using the Automatic1111 SD backend for the AI-Render plugin. Follow instructions [here](https://github.com/benrugg/AI-Render/wiki/Local-Installation) to setup local SD backend.
+
+Your AI-Render preferences should be configured as shown; the highlighted part should match your terminal output:
+![image](https://user-images.githubusercontent.com/87458719/231390322-59a54a09-520a-4a08-b658-6e37bd63e932.png)
+
+
+The [AI-Render README](https://github.com/benrugg/AI-Render/blob/main/README.md) has more details on installation and usage, as well as video tutorials.
+
+## Using AI-Render + SHARK in your Blender project
+
+- In the Render Properties tab, in the AI-Render dropdown, enable AI-Render.
+
+![image](https://user-images.githubusercontent.com/87458719/231392843-9bd51744-3ce2-464e-843a-0c4d4c96df0c.png)
+
+- Select an image size (it's usually better to upscale later than go high on the img2img resolution here.)
+
+![image](https://user-images.githubusercontent.com/87458719/231394288-0c4ab8c5-dc30-4dbe-8bc1-7520ded5efe8.png)
+
+- From here, you can enter a prompt and configure img2img Stable Diffusion parameters, and AI-Render will run SHARK SD img2img on the rendered scene.
+- AI-Render has useful presets for aesthetic styles, so you should be able to keep your subject prompt simple and focus on creating a decent Blender scene to start from.
+
+![image](https://user-images.githubusercontent.com/87458719/231440729-2fe69586-41cb-4274-9ce7-f6c08def600b.png)
+
+## Examples:
+Scene (Input image):
+
+![blender-sample-2](https://user-images.githubusercontent.com/87458719/231450408-0e680086-3e52-4962-a5c1-c703a94d1583.png)
+
+Prompt:
+"A bowl of tangerines in front of rocks, masterpiece, oil on canvas, by Georgia O'Keefe, trending on artstation, landscape painting by Caspar David Friedrich"
+
+Negative Prompt (default):
+"ugly, bad art, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, tiling, signature, cut off, draft"
+
+Example output:
+
+![blender-sample-2_out](https://user-images.githubusercontent.com/87458719/231451145-a0b56897-a7d0-4add-bbed-7e8af21a65df.png)
+
+
+
+
+
+
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,7 +16,7 @@ parameterized

 # Add transformers, diffusers and scipy since it most commonly used
 transformers
-diffusers @ git+https://github.com/huggingface/diffusers@main
+diffusers @ git+https://github.com/huggingface/diffusers@e47459c80f6f6a5a1c19d32c3fd74edf94f47aa2
 scipy
 ftfy
 gradio
--- a/shark/shark_downloader.py
+++ b/shark/shark_downloader.py
@@ -202,7 +202,7 @@ def download_model(
    dyn_str = "_dynamic" if dynamic else ""
    os.makedirs(WORKDIR, exist_ok=True)
    shark_args.shark_prefix = get_sharktank_prefix()
-    if import_args["batch_size"] != 1:
+    if import_args["batch_size"] and import_args["batch_size"] != 1:
        model_dir_name = (
            model_name
            + "_"
Author	SHA1	Message	Date
m68k-fr	35de7e27fa	[Web] remove txt2img ui dependencies from png import metadata (#1275 )	2023-04-12 07:32:47 -10:00
yzhang93	467f900759	Add auto-tuner to SD apps (#1291 )	2023-04-12 09:21:17 -07:00
Ean Garvey	0bd9d582c7	Add documentation for using SHARK with AI-Render (#1296 )	2023-04-12 03:09:34 -10:00
jinchen62	428cfe8dae	Fix low vram mode issues (#1295 ) - add ondemand back to img2img - workaround memory leak for batch count	2023-04-11 17:59:09 -07:00
Ean Garvey	f17915bedc	Fix batch size appending to model name. (#1294 ) * Update shark_downloader.py * Update shark_downloader.py	2023-04-11 15:34:25 -05:00
Gaurav Shukla	1b49b5149a	[SD] Add Img2Img rest API Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>	2023-04-11 23:06:58 +05:30
jinchen62	3002793301	Unload clip on demand and workaround memory leak (#1283 )	2023-04-10 16:59:03 -07:00
Phaneesh Barwaria	d25ef5529f	Add fix for vae fp32 Upscalar (#1284 ) - fixes size mismatch error for upscalar vae	2023-04-07 14:36:40 -05:00
Ean Garvey	308856a947	Touch unet if base cfg needed for SD pipeline init (#1281 )	2023-04-05 03:02:29 -05:00
m68k-fr	151b4e142f	[SD] Fix encoder error for model_max_length not beeing 77 (#1278 ) Co-authored-by: powderluv <powderluv@users.noreply.github.com>	2023-04-04 22:39:29 -07:00
Ean Garvey	e5a69a7c36	pin diffusers to e47459c (#1279 )	2023-04-04 18:29:21 -07:00