From acfa5617d747a50cc4c81be9b8140345bd70a24f Mon Sep 17 00:00:00 2001 From: George Petterson Date: Thu, 26 Sep 2024 11:12:44 -0500 Subject: [PATCH] 1.5 support --- apps/shark_studio/api/shark_api.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/apps/shark_studio/api/shark_api.py b/apps/shark_studio/api/shark_api.py index b171bb3d3a..ed13b4f697 100644 --- a/apps/shark_studio/api/shark_api.py +++ b/apps/shark_studio/api/shark_api.py @@ -95,11 +95,9 @@ def generate_images( (existing_args, pipeline) = pipelines[base_model] if not existing_args or not pipeline or not pipeline_args == existing_args: # TODO: Initialize new pipeline - if base_model == "sd1.5": - pass - elif base_model == "sd2": + if base_model in ["sd1.5", "sd2"]: new_pipeline = SharkSDPipeline( - hf_model_name="stabilityai/stable-diffusion-2-1", + hf_model_name=("stabilityai/stable-diffusion-2-1" if base_model == "sd2" else "stabilityai/stable-diffusion-1-5"), scheduler_id=scheduler, height=height, width=width,