From 5d239bd76f7597877d48baa5e7e79d20e885e21f Mon Sep 17 00:00:00 2001
From: db0 <mail@dbzer0.com>
Date: Thu, 26 Sep 2024 22:02:37 +0200
Subject: [PATCH] fix: correct report of censor:

---
 .../process_management/process_manager.py     | 32 ++++++++++---------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/horde_worker_regen/process_management/process_manager.py b/horde_worker_regen/process_management/process_manager.py
index 1f67fa44..78043b4d 100644
--- a/horde_worker_regen/process_management/process_manager.py
+++ b/horde_worker_regen/process_management/process_manager.py
@@ -1806,23 +1806,25 @@ def receive_and_handle_process_messages(self) -> None:
                     completed_job_info.state = GENERATION_STATE.faulted
                 elif num_images_censored > 0:
                     completed_job_info.censored = True
-                    if num_images_csam > 0:
-                        new_meta_entry = GenMetadataEntry(
-                            type=METADATA_TYPE.censorship,
-                            value=METADATA_VALUE.csam,
-                        )
-                        completed_job_info.job_image_results[i].generation_faults.append(new_meta_entry)
-                        completed_job_info.state = GENERATION_STATE.csam
-                    else:
-                        new_meta_entry = GenMetadataEntry(
-                            type=METADATA_TYPE.censorship,
-                            value=METADATA_VALUE.nsfw,
-                        )
-                        completed_job_info.job_image_results[i].generation_faults.append(new_meta_entry)
-                        completed_job_info.state = GENERATION_STATE.censored
+                    for i in range(len(completed_job_info.job_image_results)):
+                        if message.safety_evaluations[i].is_csam:
+                            new_meta_entry = GenMetadataEntry(
+                                type=METADATA_TYPE.censorship,
+                                value=METADATA_VALUE.csam,
+                            )
+                            completed_job_info.job_image_results[i].generation_faults.append(new_meta_entry)
+                            completed_job_info.state = GENERATION_STATE.csam
+                        elif message.safety_evaluations[i].is_nsfw:
+                            new_meta_entry = GenMetadataEntry(
+                                type=METADATA_TYPE.censorship,
+                                value=METADATA_VALUE.nsfw,
+                            )
+                            completed_job_info.job_image_results[i].generation_faults.append(new_meta_entry)
+                            if completed_job_info.state != GENERATION_STATE.csam:
+                                completed_job_info.state = GENERATION_STATE.censored
                 else:
                     completed_job_info.censored = False
-
+                # logger.debug([c.generation_faults for c in completed_job_info.job_image_results])
                 self.completed_jobs.append(completed_job_info)
 
     def preload_models(self) -> bool: