Small fixes to igpu, SDXL-turbo

nod-ai · May 30, 2024 · c7bd54d · c7bd54d
1 parent 1adfb3b
commit c7bd54d
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 85 deletions.
diff --git 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
@@ -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
@@ -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"),
 }