warm up model

Dockerfile

@@ -3,6 +3,9 @@ FROM python:3.8
 COPY requirements.txt requirements.txt
 RUN pip3 install -r requirements.txt
 
+ENV TRANSFORMERS_CACHE /var/cache/transformers
+RUN transformers-cli download CompVis/stable-diffusion-v1-4
+
 WORKDIR $HOME/app
 
 COPY . $HOME/app

img_nsfw.py

@@ -11,12 +11,13 @@ 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", 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)
@@ -78,10 +79,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():