Multi-head self-attention with optional 2D RoPE. x: (B, N, C) rope: VisionRotaryEmbeddingFast or None rope_ids: (B, N) or (N,

Forward pass of SiT with SPRINT. x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) t: (N,

Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes.

Get the directory path of the module containing the given object name.

Return the fully-qualified name of a top-level function.

r"""Register an import hook that is called whenever a persistent object is being unpickled. A typical use case is to patch the pickled source

r"""Initializes `torch_utils.training_stats` for collecting statistics across multiple processes. This function must be called after `tor

List all files recursively in a given directory while ignoring given file and directory names. Returns list of tuples containing both absolute and

Download the given URL and return a binary-mode file object to access the data.

Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. These may be fra

r"""Broadcasts the given set of scalars by the first process (`rank = 0`), but ignores any scalars provided by the other processes. See `repor

Take the mean over all non-batch dimensions.

Worker function for parallel image transformation.

Calculate the product of the tuple elements.