watermark optional (replicate#1)

* watermark optional

Author: daanelson

cog.yaml

@@ -11,7 +11,7 @@ build:
     - "libsm6"
     - "libxext6"
   python_packages:
-    - "git+https://github.com/huggingface/diffusers.git@34abee090750d22fef357bfb1ffd564c961b9e1d"
+    - "diffusers==0.19.3"
     - "torch==2.0.1"
     - "transformers==4.31.0"
     - "invisible-watermark==0.2.0"

predict.py

@@ -249,6 +249,10 @@ def predict(
             description="for base_image_refiner, the number of steps to refine, defaults to num_inference_steps",
             default=None,
         ),
+        apply_watermark: bool = Input(
+            description="Applies a watermark to enable determining if an image is generated in downstream applications. If you have other provisions for generating or deploying images safely, you can use this to disable watermarking.",
+            default=True
+        )
     ) -> List[Path]:
         """Run a single prediction on the model"""
         if seed is None:
@@ -284,6 +288,12 @@ def predict(
         elif refine == "base_image_refiner":
             sdxl_kwargs["output_type"] = "latent"
 
+        if not apply_watermark:
+            # toggles watermark for this prediction
+            watermark_cache = pipe.watermark
+            pipe.watermark = None
+            self.refiner.watermark = None
+
         pipe.scheduler = SCHEDULERS[scheduler].from_config(pipe.scheduler.config)
         generator = torch.Generator("cuda").manual_seed(seed)
 
@@ -307,7 +317,11 @@ def predict(
                 common_args["num_inference_steps"] = refine_steps
 
             output = self.refiner(**common_args, **refiner_kwargs)
-
+        
+        if not apply_watermark:
+            pipe.watermark = watermark_cache
+            self.refiner.watermark = watermark_cache
+        
         _, has_nsfw_content = self.run_safety_checker(output.images)
 
         output_paths = []