From b658d58c2052aa852ea67a4a67d00374d698749b Mon Sep 17 00:00:00 2001
From: zsxkib <zsakib.ahamed@gmail.com>
Date: Thu, 30 May 2024 12:04:11 +0000
Subject: [PATCH] Fix image watermark and quality issues in Instant ID model

---
 .dockerignore  |  7 ++++---
 .gitignore     | 26 ++++++++++++++++++++++----
 cog/predict.py | 24 ++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index 9a406309..b848ba91 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -3,9 +3,10 @@
 # https://docs.docker.com/engine/reference/builder/#dockerignore-file
 
 # Replicate
-checkpoints/
-models/
-safety-cache/
+/checkpoints/
+/models/
+/safety-cache/
+/gradio_cached_examples/
 *.mp4
 *.pth
 *.pt
diff --git a/.gitignore b/.gitignore
index 19c8519d..32361b8c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -163,10 +163,28 @@ checkpoints/
 models/
 
 # Cog
-.cog/
-checkpoints/
-models/
-safety-cache/
+/.cog/
+/checkpoints/
+/models/
+/safety-cache/
 *.tar
 .vscode
 gradio_cached_examples
+cog/test_batchsize.py
+input.png
+output_*.png
+output.*.png
+output_image_*.png
+output_image.*.png
+output_*.webp
+output.*.webp
+output_image_*.webp
+output_image.*.webp
+output_*.jpg
+output.*.jpg
+output_image_*.jpg
+output_image.*.jpg
+output_*.jpeg
+output.*.jpeg
+output_image_*.jpeg
+output_image.*.jpeg
\ No newline at end of file
diff --git a/cog/predict.py b/cog/predict.py
index 61f343b0..9432b4a3 100644
--- a/cog/predict.py
+++ b/cog/predict.py
@@ -33,7 +33,7 @@
 from diffusers.pipelines.stable_diffusion.safety_checker import (
     StableDiffusionSafetyChecker,
 )
-from pipeline_stable_diffusion_xl_instantid import (
+from pipeline_stable_diffusion_xl_instantid_full import (
     StableDiffusionXLInstantIDPipeline,
     draw_kps,
 )
@@ -214,12 +214,13 @@ def setup(self) -> None:
         if not os.path.exists(MODELS_CACHE):
             download_weights(MODELS_URL, MODELS_CACHE)
 
+        self.face_detection_input_width, self.face_detection_input_height = 640, 640
         self.app = FaceAnalysis(
             name="antelopev2",
             root="./",
             providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
         )
-        self.app.prepare(ctx_id=0, det_size=(640, 640))
+        self.app.prepare(ctx_id=0, det_size=(self.face_detection_input_width, self.face_detection_input_height))
 
         # Path to InstantID models
         self.face_adapter = f"./checkpoints/ip-adapter.bin"
@@ -548,6 +549,18 @@ def predict(
                 "RealVisXL_V4.0_Lightning",
             ],
         ),
+        face_detection_input_width: int = Input(
+            description="Width of the input image for face detection",
+            default=640,
+            ge=640,
+            le=4096,
+        ),
+        face_detection_input_height: int = Input(
+            description="Height of the input image for face detection",
+            default=640,
+            ge=640,
+            le=4096,
+        ),
         scheduler: str = Input(
             description="Scheduler",
             choices=[
@@ -670,6 +683,13 @@ def predict(
         if sdxl_weights != self.base_weights:
             self.load_weights(sdxl_weights)
 
+        # Resize the output if the provided dimensions are different from the current ones
+        if self.face_detection_input_width != face_detection_input_width or self.face_detection_input_height != face_detection_input_height:
+            print(f"[!] Resizing output to {face_detection_input_width}x{face_detection_input_height}")
+            self.face_detection_input_width = face_detection_input_width
+            self.face_detection_input_height = face_detection_input_height
+            self.app.prepare(ctx_id=0, det_size=(self.face_detection_input_width, self.face_detection_input_height))
+
         # Set up ControlNet selection and their respective strength values (if any)
         controlnet_selection = []
         if pose_strength > 0 and enable_pose_controlnet: