Spaces:
Running on L40S
Running on L40S
File size: 8,864 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import contextlib
import dataclasses
import traceback
from abc import ABC, abstractmethod
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, ContextManager, Self, Sequence, final
import torch
import torch.profiler
from cosmos_framework.inference.common.args import GuardrailArgs, SampleArgs, SampleOutputs, SetupArgs
from cosmos_framework.inference.common.init import is_rank0
from cosmos_framework.utils import log
from cosmos_framework.utils.misc import TrainingTimer
if TYPE_CHECKING:
from cosmos_framework.auxiliary.guardrail.common.core import GuardrailRunner
@contextlib.contextmanager
def sync_distributed_errors():
"""Catches local exceptions and synchronizes the error state across all distributed ranks.
Raises a DistributedError on all ranks if ANY rank encountered an exception.
"""
error_flag = torch.zeros(1, dtype=torch.int32, device="cuda") # [1]
local_error: Exception | None = None
try:
yield
except Exception as e:
error_flag += 1
local_error = e
if torch.distributed.is_initialized():
# Sync the error count across all GPUs
torch.distributed.all_reduce(error_flag, op=torch.distributed.ReduceOp.SUM)
if error_flag.item() > 0:
# If we got here, somebody failed.
# Ranks that failed will raise their actual error.
# Ranks that succeeded will raise a generic error so they gracefully abort too.
err_to_raise = local_error if local_error else RuntimeError("A different GPU rank failed.")
raise err_to_raise
@dataclass
class GuardrailRunners:
text: "GuardrailRunner"
video: "GuardrailRunner"
@classmethod
def create(cls, args: GuardrailArgs, /) -> Self:
from cosmos_framework.auxiliary.guardrail.common import presets
return cls(
text=presets.create_text_guardrail_runner(offload_model_to_cpu=args.offload_guardrail_models),
video=presets.create_video_guardrail_runner(offload_model_to_cpu=args.offload_guardrail_models),
)
@dataclass(kw_only=True)
class Inference(ABC):
"""Inference pipeline base class."""
setup_args: SetupArgs
model: torch.nn.Module
guardrails: GuardrailRunners | None
_timer: TrainingTimer | None
_timer_context: list[str] = dataclasses.field(default_factory=list)
@property
@abstractmethod
def model_config(self) -> Any:
"""Get model config."""
@classmethod
@abstractmethod
def _create(cls, setup_args: SetupArgs, /, **kwargs: Any) -> Self:
"""Create instance."""
@abstractmethod
def create_batches(
self, sample_args_list: Sequence[SampleArgs]
) -> Iterator[tuple[list[SampleArgs], dict[str, Any]]]:
"""Create batches of sample data."""
@abstractmethod
def generate_batch(
self, sample_args_list: Sequence[SampleArgs], data_batch: dict[str, Any], *, warmup: bool = False
) -> list[SampleOutputs]:
"""Generate a batch of samples."""
@final
@classmethod
def create(cls, setup_args: SetupArgs, /) -> Self:
"""Create instance."""
timer = TrainingTimer() if setup_args.benchmark else None
guardrails = GuardrailRunners.create(setup_args) if setup_args.guardrails else None
return cls._create(setup_args, guardrails=guardrails, _timer=timer)
@torch.no_grad()
@final
def generate(self, sample_args_list: list[SampleArgs]) -> list[SampleOutputs]:
"""Generate a list of samples."""
# Create batches
try:
with sync_distributed_errors():
batches = self.create_batches(sample_args_list)
except Exception as e:
return [self._handle_sample_exception(sample_args, e) for sample_args in sample_args_list]
# Generate batches
sample_outputs: list[SampleOutputs] = []
for i_batch, (sample_args_batch, data_batch) in enumerate(batches):
log.debug(f"[{i_batch + 1}] Processing batch", rank0_only=False)
if self.setup_args.profile:
profiler = torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True,
with_stack=True,
)
else:
profiler = contextlib.nullcontext()
with self._get_timer_context("warmup"):
for _ in range(self.setup_args.warmup):
with self._get_timer(f"{self.__class__.__name__}.generate_batch"):
self.generate_batch(sample_args_batch, data_batch, warmup=True)
with self._get_timer(f"{self.__class__.__name__}.generate_batch"), profiler:
sample_outputs.extend(self.generate_batch(sample_args_batch, data_batch))
if self.setup_args.profile and is_rank0():
assert isinstance(profiler, torch.profiler.profile)
sample_args = sample_args_batch[0]
profile_file = sample_args.output_dir / "profile.json.gz"
profiler.export_chrome_trace(str(profile_file))
log.success(f"Saved profile to '{profile_file}'")
return sample_outputs
def _get_timer(self, func_name: str) -> ContextManager:
if self._timer is None:
return nullcontext()
if self._timer_context:
context = ".".join(self._timer_context)
func_name = f"[{context}] {func_name}"
return self._timer(func_name)
@contextmanager
def _get_timer_context(self, func_name: str):
self._timer_context.append(func_name)
try:
yield
finally:
self._timer_context.pop()
def get_timer_results(self) -> dict | None:
if self._timer is None:
return None
return {
"all": self._timer.results,
"average": self._timer.compute_average_results(),
}
def _handle_sample_exception(self, sample_args: SampleArgs, e: Exception) -> SampleOutputs:
msg = f"Error generating sample '{sample_args.name}': {e}"
if not self.setup_args.keep_going:
raise ValueError(msg) from e
log.error(msg)
return SampleOutputs(
args=sample_args.model_dump(mode="json"), status="error", message=msg, stack_trace=traceback.format_exc()
)
@final
def _run_text_guardrail(self, name: str, prompt: str) -> None:
"""Run guardrail checks on the prompt."""
if self.guardrails is None:
return
from cosmos_framework.auxiliary.guardrail.common import presets
if not presets.run_text_guardrail(prompt, self.guardrails.text):
raise ValueError(f"Guardrail blocked prompt '{name}': '{prompt}'")
@final
def _run_video_guardrail(self, name: str, video_cthw: torch.Tensor) -> torch.Tensor:
"""Run guardrail checks on the video and apply face blur."""
if self.guardrails is None:
return video_cthw
processed_video_cthw, message = _run_video_guardrail(self.guardrails.video, video_cthw)
if processed_video_cthw is None:
raise ValueError(f"Guardrail blocked video '{name}': {message}")
return processed_video_cthw
def _run_video_guardrail(
video_guardrail_runner: "GuardrailRunner", video_cthw: torch.Tensor
) -> tuple[torch.Tensor | None, str]:
"""Run video guardrail and apply face blur.
Returns a ``(video_or_none, message)`` tuple. When the guardrail blocks
the video, ``video_or_none`` is ``None`` and ``message`` contains the
underlying reason (unsafe frame ratio, categories, etc.) as produced by
:class:`GuardrailRunner.run_safety_check`.
"""
if video_cthw.ndim != 4:
raise ValueError(f"Video tensor must have 4 dimensions, got {video_cthw.shape}")
frames_thwc = (
(video_cthw * 255.0).clamp(0.0, 255.0).to(torch.uint8).permute(1, 2, 3, 0).detach().cpu().numpy()
) # [T,H,W,C]
# Inline of presets.run_video_guardrail so we can forward `message` (the helper drops it).
is_safe, message = video_guardrail_runner.run_safety_check(frames_thwc)
if not is_safe:
log.critical(f"GUARDRAIL BLOCKED: {message}")
return None, message
frames_thwc = video_guardrail_runner.postprocess(frames_thwc)
video_cthw = (torch.from_numpy(frames_thwc).float().permute(3, 0, 1, 2) / 255.0).to( # [C,T,H,W]
video_cthw.device, dtype=video_cthw.dtype
)
return video_cthw, message
|