Upload folder using huggingface_hub

inference.py

@@ -14,18 +14,12 @@ from internals.pipelines.prompt_modifier import PromptModifier
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.util.args import apply_style_args
 from internals.util.avatar import Avatar
-from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda, clear_cuda_and_gc
-from internals.util.commons import (
-    download_image,
-    pickPoses,
-    upload_image,
-    upload_images,
-)
-from internals.util.config import (
-    num_return_sequences,
-    set_configs_from_task,
-    set_root_dir,
-)
+from internals.util.cache import (auto_clear_cuda_and_gc, clear_cuda,
+                                  clear_cuda_and_gc)
+from internals.util.commons import (download_image, pickPoses, upload_image,
+                                    upload_images)
+from internals.util.config import (num_return_sequences, set_configs_from_task,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.slack import Slack
@@ -455,10 +449,6 @@ def model_fn(model_dir):
     img2img_pipe.create(text2img_pipe)
     inpainter.create(text2img_pipe)
 
-    safety_checker.apply(text2img_pipe)
-    safety_checker.apply(img2img_pipe)
-    safety_checker.apply(controlnet)
-
     print("Logs: model loaded ....")
     return
 
@@ -474,6 +464,11 @@ def predict_fn(data, pipe):
         # Set set_environment
         set_configs_from_task(task)
 
+        # Apply safety checkers based on environment
+        safety_checker.apply(text2img_pipe)
+        safety_checker.apply(img2img_pipe)
+        safety_checker.apply(controlnet)
+
         # Apply arguments
         apply_style_args(data)
 

inference2.py

@@ -7,18 +7,17 @@ from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.inpainter import InPainter
 from internals.pipelines.object_remove import ObjectRemoval
 from internals.pipelines.prompt_modifier import PromptModifier
-from internals.pipelines.remove_background import RemoveBackground, RemoveBackgroundV2
+from internals.pipelines.remove_background import (RemoveBackground,
+                                                   RemoveBackgroundV2)
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.pipelines.upscaler import Upscaler
 from internals.util.avatar import Avatar
 from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
-from internals.util.commons import construct_default_s3_url, upload_image, upload_images
-from internals.util.config import (
-    num_return_sequences,
-    set_configs_from_task,
-    set_root_dir,
-)
+from internals.util.commons import (construct_default_s3_url, upload_image,
+                                    upload_images)
+from internals.util.config import (num_return_sequences, set_configs_from_task,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.slack import Slack
 
@@ -173,8 +172,6 @@ def model_fn(model_dir):
 
     replace_background.load(upscaler, remove_background_v2)
 
-    safety_checker.apply(inpainter)
-
     print("Logs: model loaded ....")
     return
 
@@ -186,13 +183,13 @@ def predict_fn(data, pipe):
 
     FailureHandler.handle(task)
 
-    # Set set_environment
-    set_configs_from_task(task)
-
     try:
         # Set set_environment
         set_configs_from_task(task)
 
+        # Apply safety checker based on environment
+        safety_checker.apply(inpainter)
+
         # Fetch avatars
         avatar.fetch_from_network(task.get_model_id())
 

internals/pipelines/safety_checker.py

@@ -4,6 +4,7 @@ import cv2
 import numpy as np
 import torch
 import torch.nn as nn
+from scipy.ndimage.filters import gaussian_filter
 from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel
 
 from internals.pipelines.commons import AbstractPipeline
@@ -23,10 +24,17 @@ class SafetyChecker:
         ).to("cuda")
 
     def apply(self, pipeline: AbstractPipeline):
+        model = self.model if not get_nsfw_access() else None
         if hasattr(pipeline, "pipe"):
-            pipeline.pipe.safety_checker = self.model
+            pipeline.pipe.safety_checker = model
         if hasattr(pipeline, "pipe2"):
-            pipeline.pipe2.safety_checker = self.model
+            pipeline.pipe2.safety_checker = model
+
+
+def cosine_distance(image_embeds, text_embeds):
+    normalized_image_embeds = nn.functional.normalize(image_embeds)
+    normalized_text_embeds = nn.functional.normalize(text_embeds)
+    return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
 
 
 class StableDiffusionSafetyCheckerV2(PreTrainedModel):
@@ -102,7 +110,7 @@ class StableDiffusionSafetyCheckerV2(PreTrainedModel):
                 result_img["concept_scores"][concept_idx] = round(
                     concept_cos - concept_threshold + adjustment, 3
                 )
-                if result_img["concept_scores"][concept_idx] > get_nsfw_threshold():
+                if result_img["concept_scores"][concept_idx] > 0:
                     result_img["bad_concepts"].append(concept_idx)
 
             result.append(result_img)
@@ -115,11 +123,12 @@ class StableDiffusionSafetyCheckerV2(PreTrainedModel):
             if any(has_nsfw_concepts) and not get_nsfw_access():
                 if torch.is_tensor(images) or torch.is_tensor(images[0]):
                     image = images[idx].cpu().numpy().astype(np.float32)
-                    image = cv2.blur(image, (30, 30))
+                    image = gaussian_filter(image, sigma=7)
+                    # image = cv2.blur(image, (30, 30))
                     image = torch.from_numpy(image)
                     images[idx] = image
                 else:
-                    images[idx] = cv2.blur(images[idx], (30, 30))
+                    images[idx] = gaussian_filter(images[idx], sigma=7)
 
         if any(has_nsfw_concepts):
             print("NSFW")
@@ -150,13 +159,14 @@ class StableDiffusionSafetyCheckerV2(PreTrainedModel):
 
         concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment
         # concept_scores = concept_scores.round(decimals=3)
-        has_nsfw_concepts = torch.any(concept_scores > get_nsfw_threshold(), dim=1)
+        has_nsfw_concepts = torch.any(concept_scores > 0, dim=1)
 
+        # images[has_nsfw_concepts] = 0.0  # black image
         # Blur images based on NSFW score
         # -------------------------------
         if not get_nsfw_access():
             image = images[has_nsfw_concepts].cpu().numpy().astype(np.float32)
-            image = cv2.blur(image, (30, 30))
+            image = gaussian_filter(image, sigma=7)
             image = torch.from_numpy(image)
             images[has_nsfw_concepts] = image