warm up model

img_nsfw.py

@@ -11,10 +11,13 @@ def init_nsfw_pipe():
     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)
     if torch.cuda.is_available():
+        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16",
+                                                       torch_dtype=torch.float16)
         pipe = pipe.to('cuda')
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp32",
+                                                       torch_dtype=torch.float32)
 
     def cosine_distance(image_embeds, text_embeds):
         normalized_image_embeds = nn.functional.normalize(image_embeds)
@@ -79,8 +82,9 @@ def check_nsfw(img, pipe):
     if torch.cuda.is_available():
         safety_checker_input = safety_checker_input.to('cuda')
     else:
-        safety_checker_input = safety_checker_input
+        safety_checker_input = safety_checker_input.to('cpu')
+
     from torch.cuda.amp import autocast
     with autocast():
-        _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values.half(), images=img)
+        _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values, images=img)
         return nsfw_tags