upd

internal/anime_nsfw/detector.go

@@ -46,6 +46,7 @@ const requestTimeout = 60 * time.Second
 type Detector struct {
 	url       string
 	threshold float64
+	workers   int
 	client    *http.Client
 	cmd       *exec.Cmd
 	scriptDir string
@@ -59,10 +60,13 @@ type Detector struct {
 }
 
 // Start 拉起本地 python 检测服务并等它就绪。
-func Start(pythonPath string, threshold float64) (*Detector, error) {
+func Start(pythonPath string, threshold float64, workers int) (*Detector, error) {
 	if threshold <= 0 || threshold >= 1 {
 		threshold = 0.5
 	}
+	if workers <= 0 {
+		workers = 2 // 默认 2 个 worker
+	}
 	py := resolvePython(pythonPath)
 	if py == "" {
 		return nil, fmt.Errorf("未找到 python，动漫敏感识别不可用")
@@ -85,6 +89,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
 
 	d := &Detector{
 		threshold: threshold,
+		workers:   workers,
 		client:    &http.Client{Timeout: requestTimeout},
 		python:    py,
 		modelDir:  modelDir,
@@ -108,6 +113,7 @@ func (d *Detector) launch() error {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
+	cmd.Env = append(os.Environ(), fmt.Sprintf("KM_ANIME_NSFW_WORKERS=%d", d.workers))
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stdout
 	if err := cmd.Start(); err != nil {
@@ -256,10 +262,10 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, ratings An
 		Nsfw    bool    `json:"nsfw"`
 		Score   float64 `json:"score"`
 		Ratings struct {
-			General       float64 `json:"general"`
-			Sensitive     float64 `json:"sensitive"`
-			Questionable  float64 `json:"questionable"`
-			Explicit      float64 `json:"explicit"`
+			General      float64 `json:"general"`
+			Sensitive    float64 `json:"sensitive"`
+			Questionable float64 `json:"questionable"`
+			Explicit     float64 `json:"explicit"`
 		} `json:"ratings"`
 		TopTags []map[string]interface{} `json:"top_tags"`
 	}