Change script loading trigger

nod-ai · May 31, 2024 · d301ef6 · d301ef6
1 parent 7e57c83
commit d301ef6
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py
@@ -104,16 +104,15 @@ 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()
+        self.is_custom = ".py" 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:
+        elif self.is_sdxl:
             self.turbine_pipe = SharkSDXLPipeline
             self.model_map = EMPTY_SDXL_MAP
         else: