Skip to content

Commit

Permalink
[Screenshot] don't use cuda kernal for now
Browse files Browse the repository at this point in the history
  • Loading branch information
japandotorg committed Sep 19, 2024
1 parent b7de931 commit 9562f6c
Showing 1 changed file with 2 additions and 5 deletions.
7 changes: 2 additions & 5 deletions screenshot/common/filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,10 @@
import asyncio
import io
import logging
from typing import Any, Dict, Literal, Optional, cast
from typing import Any, Dict, Optional, cast

import transformers
from PIL import Image
from tensorflow import config as conf

try:
import regex as re
Expand All @@ -41,11 +40,9 @@


class Filter:
DEVICE: Literal["cpu", "cuda"] = "cuda" if conf.list_physical_devices("GPU") else "cpu"

def __init__(self) -> None:
self.model: transformers.Pipeline = transformers.pipeline(
"image-classification", model="Falconsai/nsfw_image_detection", device=self.DEVICE
"image-classification", model="Falconsai/nsfw_image_detection"
)

@staticmethod
Expand Down

0 comments on commit 9562f6c

Please sign in to comment.