Fix image watermark and quality issues in Instant ID model

.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

.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

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: