Adding new dependencies and updating safety checker

leopedroso45 · Jun 28, 2024 · 3ca0102 · 3ca0102
1 parent 8055a10
commit 3ca0102
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,4 @@
 torch
 diffusers
-omegaconf
+omegaconf
+transformers
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="sevsd",
-    version="0.1.0",
+    version="0.2.0",
     author="Leonardo Severo",
     author_email="leopedroso45@gmail.com",
     description="A Python package to make Stable Diffusion Image Generation ridiculously easy",
@@ -19,6 +19,7 @@
     install_requires=[
         "torch",
         "diffusers",
+        "transformers",
         "omegaconf",
     ],
     include_package_data=True,

diff --git a/sevsd/setup_pipeline.py b/sevsd/setup_pipeline.py
@@ -1,5 +1,7 @@
 from sevsd.setup_device import setup_device
 from diffusers import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor
 
 def setup_pipeline(pretrained_model_link_or_path, **kwargs):
     r"""
@@ -27,8 +29,8 @@ def setup_pipeline(pretrained_model_link_or_path, **kwargs):
 
     default_kwargs = {
         "use_safetensors": False,
-        "load_safety_checker": False,
-        "requires_safety_checker": False,
+        "safety_checker": None,
+        "feature_extractor": AutoFeatureExtractor.from_pretrained(pretrained_model_link_or_path),
     }
 
     if pretrained_model_link_or_path.endswith(".safetensors"):