refactor: move safety to utils

imaginairy/api/generate.py

@@ -238,7 +238,6 @@ def _generate_single_image_compvis(
     from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
     from imaginairy.enhancers.upscale_realesrgan import upscale_image
     from imaginairy.modules.midas.api import torch_image_to_depth_map
-    from imaginairy.safety import create_safety_score
     from imaginairy.samplers import SOLVER_LOOKUP
     from imaginairy.samplers.editing import CFGEditingDenoiser
     from imaginairy.schema import ControlInput, ImagineResult, MaskMode
@@ -264,6 +263,7 @@ def _generate_single_image_compvis(
         outpaint_arg_str_parse,
         prepare_image_for_outpaint,
     )
+    from imaginairy.utils.safety import create_safety_score
 
     latent_channels = 4
     downsampling_factor = 8

imaginairy/api/generate_refiners.py

@@ -35,7 +35,6 @@ def _generate_single_image(
     from imaginairy.enhancers.describe_image_blip import generate_caption
     from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
     from imaginairy.enhancers.upscale_realesrgan import upscale_image
-    from imaginairy.safety import create_safety_score
     from imaginairy.samplers import SolverName
     from imaginairy.schema import ImagineResult
     from imaginairy.utils import get_device, randn_seeded
@@ -58,6 +57,7 @@ def _generate_single_image(
         outpaint_arg_str_parse,
         prepare_image_for_outpaint,
     )
+    from imaginairy.utils.safety import create_safety_score
 
     if dtype is None:
         dtype = torch.float16 if half_mode else torch.float32

tests/test_safety.py

@@ -1,8 +1,8 @@
 import pytest
 from PIL import Image
 
-from imaginairy.safety import create_safety_score
 from imaginairy.utils import get_device
+from imaginairy.utils.safety import create_safety_score
 from tests import TESTS_FOLDER