From 4505c4549f4d322c67d079add53f32c4624c00c8 Mon Sep 17 00:00:00 2001 From: Ean Garvey <87458719+monorimet@users.noreply.github.com> Date: Thu, 30 May 2024 11:40:42 -0500 Subject: [PATCH] Force inlined weights on igpu for now, small fixes to chatbot (#2149) * Add igpu and custom triple support. * Small fixes to igpu, SDXL-turbo * custom pipe loading * formatting * Remove old nodlogo import. --- apps/shark_studio/api/sd.py | 37 ++++++- apps/shark_studio/api/utils.py | 23 ++-- apps/shark_studio/modules/schedulers.py | 140 ++++++++++++------------ apps/shark_studio/web/ui/sd.py | 3 +- 4 files changed, 117 insertions(+), 86 deletions(-) diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py index 5f27c11c71..a09055878c 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,23 @@ } +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 @@ -78,19 +97,28 @@ 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 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 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 +150,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 +317,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..1e3032e70b 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,15 @@ 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 +89,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"), } diff --git a/apps/shark_studio/web/ui/sd.py b/apps/shark_studio/web/ui/sd.py index 0698ebda5e..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, ) @@ -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(