+ add use_cuda for get_model funcs in two OPs (#389)

modelscope · Aug 19, 2024 · 3eaba7f · 3eaba7f
1 parent 7a00933
commit 3eaba7f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/data_juicer/ops/filter/video_aesthetics_filter.py b/data_juicer/ops/filter/video_aesthetics_filter.py
@@ -159,7 +159,9 @@ def compute_stats(self, sample, rank=None, context=False):
 
             if len(frame_images) > 0:
                 # compute aesthetics_scores
-                model, processor = get_model(self.model_key, rank=rank)
+                model, processor = get_model(self.model_key,
+                                             rank=rank,
+                                             use_cuda=self.use_cuda())
                 inputs = processor(images=frame_images,
                                    return_tensors='pt').to(model.device)
                 with torch.no_grad():

diff --git a/data_juicer/ops/mapper/image_diffusion_mapper.py b/data_juicer/ops/mapper/image_diffusion_mapper.py
@@ -126,7 +126,9 @@ def _real_guidance(self, caption: str, image: Image.Image, rank=None):
         canvas = image.resize((512, 512), Image.BILINEAR)
         prompt = caption
 
-        diffusion_model = get_model(model_key=self.model_key, rank=rank)
+        diffusion_model = get_model(model_key=self.model_key,
+                                    rank=rank,
+                                    use_cuda=self.use_cuda())
 
         kwargs = dict(image=canvas,
                       prompt=[prompt],