diff --git a/meme.py b/meme.py index c0d5320..dac4bfd 100644 --- a/meme.py +++ b/meme.py @@ -79,7 +79,12 @@ class HMImagePipelineLoader: CATEGORY = "hellomeme" def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1', deployment='huggingface'): dtype = torch.float16 - pipeline = HMImagePipeline.from_pretrained("songkey/stable-diffusion-v1-5") + if deployment == 'modelscope': + from modelscope import snapshot_download + sd1_5_dir = snapshot_download('songkey/stable-diffusion-v1-5') + else: + sd1_5_dir = "songkey/stable-diffusion-v1-5" + pipeline = HMImagePipeline.from_pretrained(sd1_5_dir) pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version, modelscope=deployment=='modelscope') @@ -130,7 +135,12 @@ class HMVideoPipelineLoader: def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1', deployment='huggingface'): dtype = torch.float16 - pipeline = HMVideoPipeline.from_pretrained("songkey/stable-diffusion-v1-5") + if deployment == 'modelscope': + from modelscope import snapshot_download + sd1_5_dir = snapshot_download('songkey/stable-diffusion-v1-5') + else: + sd1_5_dir = "songkey/stable-diffusion-v1-5" + pipeline = HMVideoPipeline.from_pretrained(sd1_5_dir) pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version, modelscope=deployment=='modelscope')