Add snapshot download support

modelscope · Jun 11, 2024 · 4ba8a7a · 4ba8a7a
1 parent 6ba7341
commit 4ba8a7a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/modelscope_agent/tools/modelscope_tools/story_diffusion_tool.py b/modelscope_agent/tools/modelscope_tools/story_diffusion_tool.py
@@ -561,7 +561,7 @@ def setup(self):
         }
 
         pipe = StableDiffusionXLPipeline.from_pretrained(
-            'SG161222/RealVisXL_V4.0',
+            self.cfg.get("real_vis_path", 'SG161222/RealVisXL_V4.0'),
             torch_dtype=torch.float16,
             use_safetensors=False)
 

diff --git a/modelscope_agent/tools/modelscope_tools/text_to_sound_tool.py b/modelscope_agent/tools/modelscope_tools/text_to_sound_tool.py
@@ -44,7 +44,9 @@ def setup(self):
 
         self.device = torch.device(
             'cuda' if torch.cuda.is_available() else 'cpu')
-        self.pipeline = AudioLDM2Pipeline.from_pretrained('cvssp/audioldm2')
+        self.pipeline = AudioLDM2Pipeline.from_pretrained(
+            self.cfg.get("audioldm2_path", 'cvssp/audioldm2')
+        )
         self.pipeline = self.pipeline.to(self.device)
         self.is_initialized = True