fix setup for watermarking.

sayakpaul · sayakpaul · commit 491bc9f5a640 · 2023-07-06T09:36:25.000+05:30
diff --git a/setup.py b/setup.py
@@ -117,6 +117,7 @@
     "torchvision",
     "transformers>=4.25.1",
     "urllib3<=2.0.0",
+    "invisible-watermark"
 ]
 
 # this is a lookup table with items like:
@@ -207,6 +208,7 @@ def run(self):
     "scipy",
     "torchvision",
     "transformers",
+    "invisible-watermark",
 )
 extras["torch"] = deps_list("torch", "accelerate")
 
diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
@@ -31,13 +31,19 @@
 from ...utils import (
     is_accelerate_available,
     is_accelerate_version,
+    is_invisible_watermark_available,
     logging,
     randn_tensor,
     replace_example_docstring,
 )
 from ..pipeline_utils import DiffusionPipeline
 from . import StableDiffusionXLPipelineOutput
-from .watermark import StableDiffusionXLWatermarker
+
+
+if is_invisible_watermark_available():
+    from .watermark import StableDiffusionXLWatermarker
+else:
+    raise ImportError("StableDiffusionXLWatermarker requires the `invisible-watermark` library installed.")
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
diff --git a/tests/others/test_dependencies.py b/tests/others/test_dependencies.py
@@ -34,4 +34,6 @@ def test_backend_registration(self):
                 for backend in cls_module._backends:
                     if backend == "k_diffusion":
                         backend = "k-diffusion"
+                    elif backend == "invisible_watermark":
+                        backend = "invisible-watermark"
                     assert backend in deps, f"{backend} is not in the deps table!"

Original file line number	Diff line number	Diff line change
`@@ -117,6 +117,7 @@`
`117`	`117`	`"torchvision",`
`118`	`118`	`"transformers>=4.25.1",`
`119`	`119`	`"urllib3<=2.0.0",`
	`120`	`+ "invisible-watermark"`
`120`	`121`	`]`
`121`	`122`
`122`	`123`	`# this is a lookup table with items like:`
`@@ -207,6 +208,7 @@ def run(self):`
`207`	`208`	`"scipy",`
`208`	`209`	`"torchvision",`
`209`	`210`	`"transformers",`
	`211`	`+ "invisible-watermark",`
`210`	`212`	`)`
`211`	`213`	`extras["torch"] = deps_list("torch", "accelerate")`
`212`	`214`