Skip to content

Commit

Permalink
Change script loading trigger
Browse files Browse the repository at this point in the history
  • Loading branch information
eagarvey-amd committed May 31, 2024
1 parent 7e57c83 commit d301ef6
Showing 1 changed file with 2 additions and 3 deletions.
5 changes: 2 additions & 3 deletions apps/shark_studio/api/sd.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit d301ef6

Please sign in to comment.