init scan

img_nsfw.py

@@ -3,17 +3,18 @@ concepts = ['sexual', 'nude', 'sex', '18+', 'naked', 'nsfw', 'porn', 'dick', 'va
 special_concepts = ["small girl (approximation)", "young child", "young girl"]
 
 import dbimutils
+import torch
 
 
 def init_nsfw_pipe():
-    import torch
     from diffusers import StableDiffusionPipeline
     from torch import nn
 
     # make sure you're logged in with `huggingface-cli login`
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16",
                                                    torch_dtype=torch.float16)
-    pipe = pipe.to('cuda')
+    if torch.cuda.is_available():
+        pipe = pipe.to('cuda')
 
     def cosine_distance(image_embeds, text_embeds):
         normalized_image_embeds = nn.functional.normalize(image_embeds)
@@ -74,9 +75,9 @@ def init_nsfw_pipe():
 def check_nsfw(img, pipe):
     if isinstance(img, str):
         img = dbimutils.read_img_from_url(img)
-        safety_checker_input = pipe.feature_extractor(images=img, return_tensors="pt").to("cuda")
-    else:
-        safety_checker_input = pipe.feature_extractor(images=img, return_tensors="pt").to("cuda")
+    safety_checker_input = pipe.feature_extractor(images=img, return_tensors="pt")
+    if torch.cuda.is_available():
+        safety_checker_input = safety_checker_input.to('cuda')
     from torch.cuda.amp import autocast
     with autocast():
         _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values, images=img)