refactor: resolve CLIPFeatureExtractor deprecation warning (#152)

This commit resolves a CLIPFeatureExtractor deprecation warning thrown by the NSFW check logic.
livepeer · Aug 13, 2024 · 0433c97 · 0433c97
1 parent eeef4c4
commit 0433c97
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 14 deletions.
diff --git a/runner/app/pipelines/utils/utils.py b/runner/app/pipelines/utils/utils.py
@@ -11,7 +11,7 @@
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from PIL import Image
 from torch import dtype as TorchDtype
-from transformers import CLIPFeatureExtractor
+from transformers import CLIPImageProcessor
 from typing import Dict
 
 logger = logging.getLogger(__name__)
@@ -147,7 +147,7 @@ def __init__(
         self._safety_checker = StableDiffusionSafetyChecker.from_pretrained(
             "CompVis/stable-diffusion-safety-checker"
         ).to(self.device)
-        self._feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        self._feature_extractor = CLIPImageProcessor.from_pretrained(
             "openai/clip-vit-base-patch32"
         )
 

diff --git a/runner/test_prompts.py b/runner/test_prompts.py