modify presets

base_world_generation_pipeline.py

@@ -22,7 +22,7 @@ import numpy as np
 import torch
 
 from .t5_text_encoder import CosmosT5TextEncoder
-from . import presets as guardrail_presets
+from .presets import presets as guardrail_presets
 
 
 class BaseWorldGenerationPipeline(ABC):

presets.py

@@ -25,53 +25,59 @@ from .video_content_safety_filter import VideoContentSafetyFilter
 from .log import log
 
 
-def create_text_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
-    """Create the text guardrail runner."""
-    blocklist_checkpoint_dir = os.path.join(checkpoint_dir, "blocklist")
-    aegis_checkpoint_dir = os.path.join(checkpoint_dir, "aegis")
-    return GuardrailRunner(safety_models=[Blocklist(blocklist_checkpoint_dir), Aegis(aegis_checkpoint_dir)])
-
-
-def create_video_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
-    """Create the video guardrail runner."""
-    video_filter_checkpoint_dir = os.path.join(checkpoint_dir, "video_content_safety_filter")
-    retinaface_checkpoint_path = os.path.join(checkpoint_dir, "face_blur_filter/Resnet50_Final.pth")
-    return GuardrailRunner(
-        safety_models=[VideoContentSafetyFilter(video_filter_checkpoint_dir)],
-        postprocessors=[RetinaFaceFilter(retinaface_checkpoint_path)],
-    )
-
-
-def run_text_guardrail(prompt: str, guardrail_runner: GuardrailRunner) -> bool:
-    """Run the text guardrail on the prompt, checking for content safety.
-
-    Args:
-        prompt: The text prompt.
-        guardrail_runner: The text guardrail runner.
-
-    Returns:
-        bool: Whether the prompt is safe.
-    """
-    is_safe, message = guardrail_runner.run_safety_check(prompt)
-    if not is_safe:
-        log.critical(f"GUARDRAIL BLOCKED: {message}")
-    return is_safe
-
-
-def run_video_guardrail(frames: np.ndarray, guardrail_runner: GuardrailRunner) -> np.ndarray | None:
-    """Run the video guardrail on the frames, checking for content safety and applying face blur.
-
-    Args:
-        frames: The frames of the generated video.
-        guardrail_runner: The video guardrail runner.
-
-    Returns:
-        The processed frames if safe, otherwise None.
-    """
-    is_safe, message = guardrail_runner.run_safety_check(frames)
-    if not is_safe:
-        log.critical(f"GUARDRAIL BLOCKED: {message}")
-        return None
-
-    frames = guardrail_runner.postprocess(frames)
-    return frames
+class presets():
+
+    @staticmethod
+    def create_text_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
+        """Create the text guardrail runner."""
+        blocklist_checkpoint_dir = os.path.join(checkpoint_dir, "blocklist")
+        aegis_checkpoint_dir = os.path.join(checkpoint_dir, "aegis")
+        return GuardrailRunner(safety_models=[Blocklist(blocklist_checkpoint_dir), Aegis(aegis_checkpoint_dir)])
+
+
+    @staticmethod
+    def create_video_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
+        """Create the video guardrail runner."""
+        video_filter_checkpoint_dir = os.path.join(checkpoint_dir, "video_content_safety_filter")
+        retinaface_checkpoint_path = os.path.join(checkpoint_dir, "face_blur_filter/Resnet50_Final.pth")
+        return GuardrailRunner(
+            safety_models=[VideoContentSafetyFilter(video_filter_checkpoint_dir)],
+            postprocessors=[RetinaFaceFilter(retinaface_checkpoint_path)],
+        )
+
+
+    @staticmethod
+    def run_text_guardrail(prompt: str, guardrail_runner: GuardrailRunner) -> bool:
+        """Run the text guardrail on the prompt, checking for content safety.
+
+        Args:
+            prompt: The text prompt.
+            guardrail_runner: The text guardrail runner.
+
+        Returns:
+            bool: Whether the prompt is safe.
+        """
+        is_safe, message = guardrail_runner.run_safety_check(prompt)
+        if not is_safe:
+            log.critical(f"GUARDRAIL BLOCKED: {message}")
+        return is_safe
+
+
+    @staticmethod
+    def run_video_guardrail(frames: np.ndarray, guardrail_runner: GuardrailRunner) -> np.ndarray | None:
+        """Run the video guardrail on the frames, checking for content safety and applying face blur.
+
+        Args:
+            frames: The frames of the generated video.
+            guardrail_runner: The video guardrail runner.
+
+        Returns:
+            The processed frames if safe, otherwise None.
+        """
+        is_safe, message = guardrail_runner.run_safety_check(frames)
+        if not is_safe:
+            log.critical(f"GUARDRAIL BLOCKED: {message}")
+            return None
+
+        frames = guardrail_runner.postprocess(frames)
+        return frames

text2world_prompt_upsampler_inference.py

@@ -26,7 +26,7 @@ import re
 from .model_config import create_text_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
-from . import presets as guardrail_presets
+from .presets import presets as guardrail_presets
 from .log import log
 
 

video2world_prompt_upsampler_inference.py

@@ -29,7 +29,7 @@ from PIL import Image
 from .model_config import create_vision_language_model_config
 from .ar_model import AutoRegressiveModel
 from .inference import chat_completion
-from . import presets as guardrail_presets
+from .presets import presets as guardrail_presets
 from .log import log
 from .utils_io import load_from_fileobj