warm up model

img_nsfw.py

@@ -11,13 +11,12 @@ def init_nsfw_pipe():
     from torch import nn
 
     # make sure you're logged in with `huggingface-cli login`
-    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)
+    # 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", torch_dtype=torch.float32)
 
     def cosine_distance(image_embeds, text_embeds):
         normalized_image_embeds = nn.functional.normalize(image_embeds)
@@ -79,10 +78,10 @@ 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")
-    if torch.cuda.is_available():
-        safety_checker_input = safety_checker_input.to('cuda')
-    else:
-        safety_checker_input = safety_checker_input.to('cpu')
+    # if torch.cuda.is_available():
+    #     safety_checker_input = safety_checker_input.to('cuda')
+    # else:
+    safety_checker_input = safety_checker_input.to('cpu')
 
     from torch.cuda.amp import autocast
     with autocast():