Preserve pinned official Ideogram 4 FP8 release, components, original terms and reference code
8ebf922 verified | from __future__ import annotations | |
| import io | |
| import requests | |
| HIVE_TEXT_MODERATION_URL = "https://api.thehive.ai/api/v3/hive/text-moderation" | |
| HIVE_VISUAL_MODERATION_URL = "https://api.thehive.ai/api/v3/hive/visual-moderation" | |
| # Visual model returns floats in [0, 1]; flag a class when its value is >= this. | |
| HIVE_VISUAL_THRESHOLD = 0.9 | |
| # Text model returns integer severities (0 = none, 1+ = present). Flag at >= 1. | |
| HIVE_TEXT_THRESHOLD = 1 | |
| # Visual classes we treat as safety signals. | |
| _VISUAL_SAFETY_CLASSES = { | |
| "general_nsfw", | |
| "general_suggestive", | |
| "yes_sexual_activity", | |
| "yes_sexual_intent", | |
| "yes_realistic_nsfw", | |
| "very_bloody", | |
| "yes_animal_abuse", | |
| "yes_self_harm", | |
| "yes_confederate", | |
| "yes_kkk", | |
| "yes_nazi", | |
| } | |
| # Text classes we treat as safety signals. | |
| _TEXT_SAFETY_CLASSES = { | |
| "bullying", | |
| "child_exploitation", | |
| "child_safety", | |
| "drugs", | |
| "hate", | |
| "self_harm", | |
| "self_harm_intent", | |
| "sexual", | |
| "sexual_description", | |
| "violence", | |
| "violent_description", | |
| } | |
| def _iter_classes(payload: dict): | |
| for output in payload.get("output", []): | |
| yield from output.get("classes", []) | |
| def moderate_prompt( | |
| prompt: str, | |
| api_key: str, | |
| threshold: float = HIVE_TEXT_THRESHOLD, | |
| ) -> list[tuple[str, float]]: | |
| """Run Hive Text Moderation over `prompt`. | |
| Returns (class_name, value) tuples for safety-relevant categories that | |
| scored at or above `threshold`. Text model values are integer severities | |
| (typically 0-3), so the default threshold of 1 flags any detection. | |
| """ | |
| resp = requests.post( | |
| HIVE_TEXT_MODERATION_URL, | |
| headers={"Authorization": f"Bearer {api_key}"}, | |
| json={"input": [{"text": prompt}]}, | |
| timeout=30, | |
| ) | |
| resp.raise_for_status() | |
| return [ | |
| (cls["class"], float(cls["value"])) | |
| for cls in _iter_classes(resp.json()) | |
| if float(cls["value"]) >= threshold and cls["class"] in _TEXT_SAFETY_CLASSES | |
| ] | |
| def moderate_image( | |
| image, | |
| api_key: str, | |
| threshold: float = HIVE_VISUAL_THRESHOLD, | |
| ) -> list[tuple[str, float]]: | |
| """Run Hive Visual Content Moderation over a PIL `image`. | |
| Returns (class_name, value) tuples for categories that scored at or above | |
| `threshold`. Visual model values are floats in [0, 1]. | |
| """ | |
| buf = io.BytesIO() | |
| image.save(buf, format="PNG") | |
| buf.seek(0) | |
| resp = requests.post( | |
| HIVE_VISUAL_MODERATION_URL, | |
| headers={"Authorization": f"Bearer {api_key}"}, | |
| files={"media": ("image.png", buf, "image/png")}, | |
| timeout=60, | |
| ) | |
| resp.raise_for_status() | |
| return [ | |
| (cls["class"], float(cls["value"])) | |
| for cls in _iter_classes(resp.json()) | |
| if float(cls["value"]) >= threshold and cls["class"] in _VISUAL_SAFETY_CLASSES | |
| ] | |