From 1adfb3b24b7cda5db2382fa6de6e14997fd83edb Mon Sep 17 00:00:00 2001 From: Ean Garvey Date: Wed, 29 May 2024 13:59:13 -0500 Subject: [PATCH 1/5] Add igpu and custom triple support. --- apps/shark_studio/api/sd.py | 5 ++++- apps/shark_studio/api/utils.py | 20 +++++++++++++++++++- apps/shark_studio/web/ui/sd.py | 9 +++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py index 83574d294d..5f27c11c71 100644 --- a/apps/shark_studio/api/sd.py +++ b/apps/shark_studio/api/sd.py @@ -73,6 +73,7 @@ def __init__( scheduler: str, precision: str, device: str, + target_triple: str = None, custom_vae: str = None, num_loras: int = 0, import_ir: bool = True, @@ -91,7 +92,7 @@ def __init__( self.model_map = EMPTY_SD_MAP external_weights = "safetensors" max_length = 64 - target_backend, self.rt_device, triple = parse_device(device) + target_backend, self.rt_device, triple = parse_device(device, target_triple) pipe_id_list = [ safe_name(base_model_id), str(batch_size), @@ -273,6 +274,7 @@ def shark_sd_fn( custom_vae: str, precision: str, device: str, + target_triple: str, ondemand: bool, repeatable_seeds: bool, resample_type: str, @@ -326,6 +328,7 @@ def shark_sd_fn( "batch_size": batch_size, "precision": precision, "device": device, + "target_triple": target_triple, "custom_vae": custom_vae, "num_loras": num_loras, "import_ir": import_ir, diff --git a/apps/shark_studio/api/utils.py b/apps/shark_studio/api/utils.py index 0516255d2b..b87ee0e628 100644 --- a/apps/shark_studio/api/utils.py +++ b/apps/shark_studio/api/utils.py @@ -77,6 +77,21 @@ def get_devices_by_name(driver_name): available_devices.extend(cpu_device) cpu_device = get_devices_by_name("cpu-task") available_devices.extend(cpu_device) + print(available_devices) + for idx, device_str in enumerate(available_devices): + if "AMD Radeon(TM) Graphics =>" in device_str: + igpu_id_candidates = [ + x.split("w/")[-1].split("=>")[0] + for x in available_devices + if "M Graphics" in x + ] + for igpu_name in igpu_id_candidates: + if igpu_name: + print(f"Found iGPU: {igpu_name} for {device_str}") + available_devices[idx] = device_str.replace( + "AMD Radeon(TM) Graphics", f"AMD iGPU: {igpu_name}" + ) + break return available_devices @@ -129,7 +144,7 @@ def set_iree_runtime_flags(): set_iree_vulkan_runtime_flags(flags=vulkan_runtime_flags) -def parse_device(device_str): +def parse_device(device_str, target_override=""): from shark.iree_utils.compile_utils import ( clean_device_info, get_iree_target_triple, @@ -143,6 +158,8 @@ def parse_device(device_str): else: rt_device = rt_driver + if target_override: + return target_backend, rt_device, target_override match target_backend: case "vulkan-spirv": triple = get_iree_target_triple(device_str) @@ -168,6 +185,7 @@ def get_rocm_target_chip(device_str): "MI100": "gfx908", "MI50": "gfx906", "MI60": "gfx906", + "780M": "gfx1103", } for key in rocm_chip_map: if key in device_str: diff --git a/apps/shark_studio/web/ui/sd.py b/apps/shark_studio/web/ui/sd.py index a4df173b1c..0698ebda5e 100644 --- a/apps/shark_studio/web/ui/sd.py +++ b/apps/shark_studio/web/ui/sd.py @@ -117,6 +117,7 @@ def pull_sd_configs( custom_vae, precision, device, + target_triple, ondemand, repeatable_seeds, resample_type, @@ -175,6 +176,7 @@ def load_sd_cfg(sd_json: dict, load_sd_config: str): sd_json["custom_vae"], sd_json["precision"], sd_json["device"], + sd_json["target_triple"], sd_json["ondemand"], sd_json["repeatable_seeds"], sd_json["resample_type"], @@ -253,6 +255,11 @@ def base_model_changed(base_model_id): choices=global_obj.get_device_list(), allow_custom_value=False, ) + target_triple = gr.Textbox( + elem_id="triple", + label="Architecture", + value="", + ) with gr.Row(): ondemand = gr.Checkbox( value=cmd_opts.lowvram, @@ -691,6 +698,7 @@ def base_model_changed(base_model_id): custom_vae, precision, device, + target_triple, ondemand, repeatable_seeds, resample_type, @@ -730,6 +738,7 @@ def base_model_changed(base_model_id): custom_vae, precision, device, + target_triple, ondemand, repeatable_seeds, resample_type, From c7bd54de2338502ad9647f6fd644af762a930995 Mon Sep 17 00:00:00 2001 From: Ean Garvey Date: Thu, 30 May 2024 00:54:12 -0500 Subject: [PATCH 2/5] Small fixes to igpu, SDXL-turbo --- apps/shark_studio/api/sd.py | 9 +- apps/shark_studio/api/utils.py | 22 ++-- apps/shark_studio/modules/schedulers.py | 140 ++++++++++++------------ 3 files changed, 86 insertions(+), 85 deletions(-) diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py index 5f27c11c71..e03b5e359d 100644 --- a/apps/shark_studio/api/sd.py +++ b/apps/shark_studio/api/sd.py @@ -78,6 +78,7 @@ def __init__( num_loras: int = 0, import_ir: bool = True, is_controlled: bool = False, + external_weights: str = "safetensors", ): self.precision = precision self.compiled_pipeline = False @@ -90,7 +91,6 @@ def __init__( else: self.turbine_pipe = SharkSDPipeline self.model_map = EMPTY_SD_MAP - external_weights = "safetensors" max_length = 64 target_backend, self.rt_device, triple = parse_device(device, target_triple) pipe_id_list = [ @@ -122,9 +122,12 @@ def __init__( if triple in ["gfx940", "gfx942", "gfx90a"]: decomp_attn = False attn_spec = "mfma" - elif triple in ["gfx1100", "gfx1103"]: + elif triple in ["gfx1100", "gfx1103", "gfx1150"]: decomp_attn = False attn_spec = "wmma" + if triple in ["gfx1103", "gfx1150"]: + # external weights have issues on igpu + external_weights = None elif target_backend == "llvm-cpu": decomp_attn = False @@ -286,8 +289,6 @@ def shark_sd_fn( sd_init_image = [sd_init_image] is_img2img = True if sd_init_image[0] is not None else False - print("\n[LOG] Performing Stable Diffusion Pipeline setup...") - from apps.shark_studio.modules.shared_cmd_opts import cmd_opts import apps.shark_studio.web.utils.globals as global_obj diff --git a/apps/shark_studio/api/utils.py b/apps/shark_studio/api/utils.py index b87ee0e628..79139e2312 100644 --- a/apps/shark_studio/api/utils.py +++ b/apps/shark_studio/api/utils.py @@ -52,6 +52,13 @@ def get_devices_by_name(driver_name): set_iree_runtime_flags() available_devices = [] + rocm_devices = get_devices_by_name("rocm") + available_devices.extend(rocm_devices) + cpu_device = get_devices_by_name("cpu-sync") + available_devices.extend(cpu_device) + cpu_device = get_devices_by_name("cpu-task") + available_devices.extend(cpu_device) + from shark.iree_utils.vulkan_utils import ( get_all_vulkan_devices, ) @@ -64,20 +71,14 @@ def get_devices_by_name(driver_name): id += 1 if id != 0: print(f"vulkan devices are available.") + available_devices.extend(vulkan_devices) metal_devices = get_devices_by_name("metal") available_devices.extend(metal_devices) cuda_devices = get_devices_by_name("cuda") available_devices.extend(cuda_devices) - rocm_devices = get_devices_by_name("rocm") - available_devices.extend(rocm_devices) hip_devices = get_devices_by_name("hip") available_devices.extend(hip_devices) - cpu_device = get_devices_by_name("cpu-sync") - available_devices.extend(cpu_device) - cpu_device = get_devices_by_name("cpu-task") - available_devices.extend(cpu_device) - print(available_devices) for idx, device_str in enumerate(available_devices): if "AMD Radeon(TM) Graphics =>" in device_str: igpu_id_candidates = [ @@ -87,10 +88,9 @@ def get_devices_by_name(driver_name): ] for igpu_name in igpu_id_candidates: if igpu_name: - print(f"Found iGPU: {igpu_name} for {device_str}") - available_devices[idx] = device_str.replace( - "AMD Radeon(TM) Graphics", f"AMD iGPU: {igpu_name}" - ) + available_devices[idx] = device_str.replace( + "AMD Radeon(TM) Graphics", igpu_name + ) break return available_devices diff --git a/apps/shark_studio/modules/schedulers.py b/apps/shark_studio/modules/schedulers.py index 56df8973d0..731e5707b4 100644 --- a/apps/shark_studio/modules/schedulers.py +++ b/apps/shark_studio/modules/schedulers.py @@ -24,47 +24,47 @@ def get_schedulers(model_id): model_id, subfolder="scheduler", ) - schedulers["DDPM"] = DDPMScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["KDPM2Discrete"] = KDPM2DiscreteScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["LMSDiscrete"] = LMSDiscreteScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["DDIM"] = DDIMScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["LCMScheduler"] = LCMScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["DPMSolverMultistep"] = DPMSolverMultistepScheduler.from_pretrained( - model_id, subfolder="scheduler", algorithm_type="dpmsolver" - ) - schedulers["DPMSolverMultistep++"] = DPMSolverMultistepScheduler.from_pretrained( - model_id, subfolder="scheduler", algorithm_type="dpmsolver++" - ) - schedulers["DPMSolverMultistepKarras"] = ( - DPMSolverMultistepScheduler.from_pretrained( - model_id, - subfolder="scheduler", - use_karras_sigmas=True, - ) - ) - schedulers["DPMSolverMultistepKarras++"] = ( - DPMSolverMultistepScheduler.from_pretrained( - model_id, - subfolder="scheduler", - algorithm_type="dpmsolver++", - use_karras_sigmas=True, - ) - ) + # schedulers["DDPM"] = DDPMScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["KDPM2Discrete"] = KDPM2DiscreteScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["LMSDiscrete"] = LMSDiscreteScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["DDIM"] = DDIMScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["LCMScheduler"] = LCMScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["DPMSolverMultistep"] = DPMSolverMultistepScheduler.from_pretrained( + # model_id, subfolder="scheduler", algorithm_type="dpmsolver" + # ) + # schedulers["DPMSolverMultistep++"] = DPMSolverMultistepScheduler.from_pretrained( + # model_id, subfolder="scheduler", algorithm_type="dpmsolver++" + # ) + # schedulers["DPMSolverMultistepKarras"] = ( + # DPMSolverMultistepScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # use_karras_sigmas=True, + # ) + # ) + # schedulers["DPMSolverMultistepKarras++"] = ( + # DPMSolverMultistepScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # algorithm_type="dpmsolver++", + # use_karras_sigmas=True, + # ) + # ) schedulers["EulerDiscrete"] = EulerDiscreteScheduler.from_pretrained( model_id, subfolder="scheduler", @@ -75,24 +75,24 @@ def get_schedulers(model_id): subfolder="scheduler", ) ) - schedulers["DEISMultistep"] = DEISMultistepScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["DPMSolverSinglestep"] = DPMSolverSinglestepScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - schedulers["KDPM2AncestralDiscrete"] = ( - KDPM2AncestralDiscreteScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) - ) - schedulers["HeunDiscrete"] = HeunDiscreteScheduler.from_pretrained( - model_id, - subfolder="scheduler", - ) + # schedulers["DEISMultistep"] = DEISMultistepScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["DPMSolverSinglestep"] = DPMSolverSinglestepScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # schedulers["KDPM2AncestralDiscrete"] = ( + # KDPM2AncestralDiscreteScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) + # ) + # schedulers["HeunDiscrete"] = HeunDiscreteScheduler.from_pretrained( + # model_id, + # subfolder="scheduler", + # ) return schedulers @@ -102,17 +102,17 @@ def export_scheduler_model(model): scheduler_model_map = { "PNDM": export_scheduler_model("PNDMScheduler"), - "DPMSolverSDE": export_scheduler_model("DpmSolverSDEScheduler"), + # "DPMSolverSDE": export_scheduler_model("DpmSolverSDEScheduler"), "EulerDiscrete": export_scheduler_model("EulerDiscreteScheduler"), "EulerAncestralDiscrete": export_scheduler_model("EulerAncestralDiscreteScheduler"), - "LCM": export_scheduler_model("LCMScheduler"), - "LMSDiscrete": export_scheduler_model("LMSDiscreteScheduler"), - "DDPM": export_scheduler_model("DDPMScheduler"), - "DDIM": export_scheduler_model("DDIMScheduler"), - "DPMSolverMultistep": export_scheduler_model("DPMSolverMultistepScheduler"), - "KDPM2Discrete": export_scheduler_model("KDPM2DiscreteScheduler"), - "DEISMultistep": export_scheduler_model("DEISMultistepScheduler"), - "DPMSolverSinglestep": export_scheduler_model("DPMSolverSingleStepScheduler"), - "KDPM2AncestralDiscrete": export_scheduler_model("KDPM2AncestralDiscreteScheduler"), - "HeunDiscrete": export_scheduler_model("HeunDiscreteScheduler"), + # "LCM": export_scheduler_model("LCMScheduler"), + # "LMSDiscrete": export_scheduler_model("LMSDiscreteScheduler"), + # "DDPM": export_scheduler_model("DDPMScheduler"), + # "DDIM": export_scheduler_model("DDIMScheduler"), + # "DPMSolverMultistep": export_scheduler_model("DPMSolverMultistepScheduler"), + # "KDPM2Discrete": export_scheduler_model("KDPM2DiscreteScheduler"), + # "DEISMultistep": export_scheduler_model("DEISMultistepScheduler"), + # "DPMSolverSinglestep": export_scheduler_model("DPMSolverSingleStepScheduler"), + # "KDPM2AncestralDiscrete": export_scheduler_model("KDPM2AncestralDiscreteScheduler"), + # "HeunDiscrete": export_scheduler_model("HeunDiscreteScheduler"), } From e6b4885b0ef33fb3067d4f6ca7527037b1903348 Mon Sep 17 00:00:00 2001 From: Ean Garvey Date: Thu, 30 May 2024 03:45:18 -0500 Subject: [PATCH 3/5] custom pipe loading --- apps/shark_studio/api/sd.py | 32 ++++++++++++++++++++++++++++++++ apps/shark_studio/web/ui/sd.py | 1 + 2 files changed, 33 insertions(+) diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py index e03b5e359d..e9ef69a98a 100644 --- a/apps/shark_studio/api/sd.py +++ b/apps/shark_studio/api/sd.py @@ -5,6 +5,8 @@ import json import numpy as np import copy +import importlib.util +import sys from tqdm.auto import tqdm from pathlib import Path @@ -56,6 +58,24 @@ } + +def load_script(source, module_name): + """ + reads file source and loads it as a module + + :param source: file to load + :param module_name: name of module to register in sys.modules + :return: loaded module + """ + + spec = importlib.util.spec_from_file_location(module_name, source) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + return module + + class StableDiffusion: # This class is responsible for executing image generation and creating # /managing a set of compiled modules to run Stable Diffusion. The init @@ -85,6 +105,18 @@ def __init__( self.base_model_id = base_model_id self.custom_vae = custom_vae self.is_sdxl = "xl" in self.base_model_id.lower() + self.is_custom = "custom" in self.base_model_id.lower() + if self.is_custom: + custom_module = load_script( + os.path.join( + get_checkpoints_path("scripts"), + self.base_model_id + ), + "custom_pipeline" + ) + self.turbine_pipe = custom_module.StudioPipeline + self.model_map = custom_module.MODEL_MAP + if self.is_sdxl: self.turbine_pipe = SharkSDXLPipeline self.model_map = EMPTY_SDXL_MAP diff --git a/apps/shark_studio/web/ui/sd.py b/apps/shark_studio/web/ui/sd.py index 0698ebda5e..7fe3962ae8 100644 --- a/apps/shark_studio/web/ui/sd.py +++ b/apps/shark_studio/web/ui/sd.py @@ -282,6 +282,7 @@ def base_model_changed(base_model_id): elem_id="custom_model", value="stabilityai/stable-diffusion-2-1-base", choices=sd_default_models, + allow_custom_value=True, ) # base_model_id with gr.Row(): height = gr.Slider( From d8bb4db800ef1de584b729ced5fbcbf9589bcca9 Mon Sep 17 00:00:00 2001 From: Ean Garvey Date: Thu, 30 May 2024 03:46:47 -0500 Subject: [PATCH 4/5] formatting --- apps/shark_studio/api/sd.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py index e9ef69a98a..a09055878c 100644 --- a/apps/shark_studio/api/sd.py +++ b/apps/shark_studio/api/sd.py @@ -58,7 +58,6 @@ } - def load_script(source, module_name): """ reads file source and loads it as a module @@ -108,11 +107,8 @@ def __init__( self.is_custom = "custom" in self.base_model_id.lower() if self.is_custom: custom_module = load_script( - os.path.join( - get_checkpoints_path("scripts"), - self.base_model_id - ), - "custom_pipeline" + os.path.join(get_checkpoints_path("scripts"), self.base_model_id), + "custom_pipeline", ) self.turbine_pipe = custom_module.StudioPipeline self.model_map = custom_module.MODEL_MAP From 8478106f730482c2d63f6adfdafbeb50560b789d Mon Sep 17 00:00:00 2001 From: Ean Garvey Date: Thu, 30 May 2024 11:37:20 -0500 Subject: [PATCH 5/5] Remove old nodlogo import. --- apps/shark_studio/web/ui/sd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/shark_studio/web/ui/sd.py b/apps/shark_studio/web/ui/sd.py index 7fe3962ae8..ee8bf77f58 100644 --- a/apps/shark_studio/web/ui/sd.py +++ b/apps/shark_studio/web/ui/sd.py @@ -32,7 +32,7 @@ ) from apps.shark_studio.modules.shared_cmd_opts import cmd_opts from apps.shark_studio.web.ui.utils import ( - nodlogo_loc, + amdlogo_loc, none_to_str_none, str_none_to_none, )