Spaces:
Running on L40S
Running on L40S
File size: 6,561 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import os
from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import Optional
import torch
from cosmos_framework.utils.env_parsers.cred_env_parser import CRED_ENVS
class VideoTokenizerInterface(ABC):
def __init__(self, object_store_credential_path_pretrained: Optional[str] = None):
assert object_store_credential_path_pretrained is None or isinstance(
object_store_credential_path_pretrained, str
)
if object_store_credential_path_pretrained is None:
self.backend_args = None
elif os.path.exists(object_store_credential_path_pretrained) or CRED_ENVS.APP_ENV in ["prod", "dev", "stg"]:
self.backend_args = {
"backend": "s3",
"path_mapping": None,
"s3_credential_path": object_store_credential_path_pretrained,
}
else:
raise FileNotFoundError(
f"Invalid object_store_credential_path_pretrained: {object_store_credential_path_pretrained} and APP_ENV is not prod/dev/stg"
)
@abstractmethod
def reset_dtype(self):
"""
Reset the dtype of the model to the dtype its weights were trained with or quantized to.
"""
pass
@abstractmethod
def encode(self, state: torch.Tensor) -> torch.Tensor:
pass
@abstractmethod
def decode(self, latent: torch.Tensor) -> torch.Tensor:
pass
@abstractmethod
def get_latent_num_frames(self, num_pixel_frames: int) -> int:
pass
@abstractmethod
def get_pixel_num_frames(self, num_latent_frames: int) -> int:
pass
@property
@abstractmethod
def spatial_compression_factor(self) -> int:
pass
@property
@abstractmethod
def temporal_compression_factor(self) -> int:
pass
@property
@abstractmethod
def spatial_resolution(self) -> int:
pass
@property
@abstractmethod
def pixel_chunk_duration(self):
pass
@property
@abstractmethod
def latent_chunk_duration(self):
pass
@property
@abstractmethod
def latent_ch(self) -> int:
pass
def compile_encode(
self,
warmup_resolutions: Sequence[str],
output_dir: str,
aspect_ratio: str | None = None,
) -> None:
"""AOT-compile the tokenizer for the given resolutions.
Subclasses that support AOT compilation should override this method.
The default raises ``NotImplementedError``.
Args:
warmup_resolutions: Resolution keys to compile for.
output_dir: Root directory where compiled artifacts are stored
(typically ``config.job.path_local``).
aspect_ratio: If given, only compile this single aspect ratio.
"""
raise NotImplementedError(f"{type(self).__name__} does not support compilation")
@property
def is_chunk_overlap(self):
return False
@property
def is_causal(self):
return True
class AudioTokenizerInterface(ABC):
"""Abstract interface for audio tokenizers."""
def __init__(self, object_store_credential_path_pretrained: Optional[str] = None):
assert object_store_credential_path_pretrained is None or isinstance(
object_store_credential_path_pretrained, str
)
if not object_store_credential_path_pretrained:
self.backend_args = None
elif os.path.exists(object_store_credential_path_pretrained) or CRED_ENVS.APP_ENV in ["prod", "dev", "stg"]:
self.backend_args = {
"backend": "s3",
"path_mapping": None,
"s3_credential_path": object_store_credential_path_pretrained,
}
else:
raise FileNotFoundError(
f"Invalid object_store_credential_path_pretrained: {object_store_credential_path_pretrained} and APP_ENV is not prod/dev/stg"
)
@abstractmethod
def reset_dtype(self):
"""
Reset the dtype of the model to the dtype its weights were trained with or quantized to.
"""
pass
@abstractmethod
def encode(self, audio: torch.Tensor, force_pad: bool = False) -> torch.Tensor:
"""
Encode audio waveform to latent representation.
Args:
audio: Input audio tensor of shape [B, C, T] where:
B = batch size, C = audio channels, T = time samples
force_pad: Whether to force padding to match compression factor
Returns:
Latent tensor of shape [B, latent_ch, T']
"""
pass
@abstractmethod
def decode(self, latent: torch.Tensor) -> torch.Tensor:
"""
Decode latent representation to audio waveform.
Args:
latent: Latent tensor of shape [B, latent_ch, T']
Returns:
Audio tensor of shape [B, C, T]
"""
pass
@abstractmethod
def get_latent_num_samples(self, num_audio_samples: int) -> int:
"""
Calculate the number of latent time samples from audio samples.
Args:
num_audio_samples: Number of audio samples
Returns:
Number of latent time samples
"""
pass
@abstractmethod
def get_audio_num_samples(self, num_latent_samples: int) -> int:
"""
Calculate the number of audio samples from latent samples.
Args:
num_latent_samples: Number of latent time samples
Returns:
Number of audio samples
"""
pass
@property
@abstractmethod
def temporal_compression_factor(self) -> int:
"""
Temporal compression factor (downsampling ratio).
audio_samples = latent_samples * temporal_compression_factor
"""
pass
@property
@abstractmethod
def sample_rate(self) -> int:
"""Audio sample rate in Hz."""
pass
@property
@abstractmethod
def audio_channels(self) -> int:
"""Number of audio channels (e.g., 1 for mono, 2 for stereo)."""
pass
@property
@abstractmethod
def latent_ch(self) -> int:
"""Number of latent channels."""
pass
@property
def is_causal(self) -> bool:
"""Whether the model is causal (for streaming)."""
return False
|