SYMBOL INDEX (125 symbols across 11 files)

FILE: batch_inference.py
  function read_prompts (line 14) | def read_prompts(path: str) -> list[str]:
  function slugify (line 30) | def slugify(text: str, max_len: int = 60) -> str:
  function select_device (line 38) | def select_device() -> str:
  function main (line 58) | def main():

FILE: inference.py
  function main (line 14) | def main():

FILE: src/tools/generate_manifest.py
  function compute_md5 (line 15) | def compute_md5(file_path: Path, chunk_size: int = 8192) -> str:
  function get_essential_files (line 24) | def get_essential_files(model_dir: Path) -> List[Path]:
  function main (line 51) | def main():

FILE: src/utils/attention.py
  class AttentionBackend (line 51) | class AttentionBackend(str, Enum):
    method print_available_backends (line 67) | def print_available_backends(cls):
  function register_backend (line 77) | def register_backend(name: str, constraints: Optional[List[Callable]] = ...
  function _check_device_cuda (line 87) | def _check_device_cuda(query: torch.Tensor, **kwargs) -> None:
  function _check_qkv_dtype_bf16_or_fp16 (line 92) | def _check_qkv_dtype_bf16_or_fp16(query: torch.Tensor, **kwargs) -> None:
  function _check_device_mps (line 97) | def _check_device_mps(query: torch.Tensor, **kwargs) -> None:
  function _process_mask (line 102) | def _process_mask(attn_mask: Optional[torch.Tensor], dtype: torch.dtype):
  function _normalize_attn_mask (line 119) | def _normalize_attn_mask(attn_mask: torch.Tensor, batch_size: int, seq_l...
  function _prepare_for_flash_attn_varlen_without_mask (line 148) | def _prepare_for_flash_attn_varlen_without_mask(
  function _prepare_for_flash_attn_varlen_with_mask (line 164) | def _prepare_for_flash_attn_varlen_with_mask(
  function _prepare_for_flash_attn_varlen (line 184) | def _prepare_for_flash_attn_varlen(
  function _flash_attention (line 197) | def _flash_attention(
  function _flash_varlen_attention (line 223) | def _flash_varlen_attention(
  function _flash_attention_3 (line 277) | def _flash_attention_3(
  function _flash_varlen_attention_3 (line 309) | def _flash_varlen_attention_3(
  function _mps_flash_attention (line 373) | def _mps_flash_attention(
  function _native_attention_wrapper (line 410) | def _native_attention_wrapper(
  function _native_flash_attention (line 440) | def _native_flash_attention(
  function _math_attention (line 462) | def _math_attention(*args, **kwargs):
  function _native_attention (line 467) | def _native_attention(*args, **kwargs):
  function dispatch_attention (line 471) | def dispatch_attention(
  function set_attention_backend (line 508) | def set_attention_backend(backend: Union[str, AttentionBackend, None]):

FILE: src/utils/helpers.py
  function format_bytes (line 14) | def format_bytes(size: float) -> str:
  function print_memory_stats (line 28) | def print_memory_stats(stage: str) -> None:
  function compute_file_md5 (line 52) | def compute_file_md5(file_path: Path, chunk_size: int = 8192) -> str:
  function load_manifest (line 61) | def load_manifest(manifest_file: Path) -> Dict[str, Optional[str]]:
  function verify_file_integrity (line 94) | def verify_file_integrity(
  function ensure_model_weights (line 135) | def ensure_model_weights(

FILE: src/utils/import_utils.py
  function is_flash_attn_available (line 6) | def is_flash_attn_available():
  function is_flash_attn_3_available (line 10) | def is_flash_attn_3_available():
  function is_torch_version (line 14) | def is_torch_version(operator: str, version: str):

FILE: src/utils/loader.py
  function load_config (line 45) | def load_config(config_path: str) -> dict:
  function load_sharded_safetensors (line 50) | def load_sharded_safetensors(weight_dir: Path, device: str = "cuda", dty...
  function load_from_local_dir (line 80) | def load_from_local_dir(

FILE: src/zimage/autoencoder.py
  class AutoencoderKLOutput (line 12) | class AutoencoderKLOutput:
  class AutoencoderConfig (line 16) | class AutoencoderConfig:
    method __init__ (line 17) | def __init__(self, **kwargs):
    method get (line 20) | def get(self, key, default=None):
    method __getattr__ (line 23) | def __getattr__(self, name):
  function swish (line 27) | def swish(x):
  class ResnetBlock2D (line 31) | class ResnetBlock2D(nn.Module):
    method __init__ (line 32) | def __init__(self, in_channels, out_channels=None, dropout=0.0, temb_c...
    method forward (line 51) | def forward(self, input_tensor, temb=None):
  class Attention (line 69) | class Attention(nn.Module):
    method __init__ (line 70) | def __init__(self, in_channels, heads=1, dim_head=None, groups=32, eps...
    method forward (line 81) | def forward(self, hidden_states):
  class Downsample2D (line 101) | class Downsample2D(nn.Module):
    method __init__ (line 102) | def __init__(self, channels, with_conv=True, out_channels=None, paddin...
    method forward (line 109) | def forward(self, hidden_states):
  class Upsample2D (line 116) | class Upsample2D(nn.Module):
    method __init__ (line 117) | def __init__(self, channels, with_conv=True, out_channels=None):
    method forward (line 124) | def forward(self, hidden_states):
  class DownEncoderBlock2D (line 131) | class DownEncoderBlock2D(nn.Module):
    method __init__ (line 132) | def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps...
    method forward (line 147) | def forward(self, hidden_states):
  class UpDecoderBlock2D (line 160) | class UpDecoderBlock2D(nn.Module):
    method __init__ (line 161) | def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps...
    method forward (line 174) | def forward(self, hidden_states):
  class UNetMidBlock2D (line 185) | class UNetMidBlock2D(nn.Module):
    method __init__ (line 186) | def __init__(self, in_channels, resnet_eps=1e-6, resnet_groups=32, att...
    method forward (line 196) | def forward(self, hidden_states):
  class Encoder (line 204) | class Encoder(nn.Module):
    method __init__ (line 205) | def __init__(
    method forward (line 244) | def forward(self, x):
  class Decoder (line 255) | class Decoder(nn.Module):
    method __init__ (line 256) | def __init__(
    method forward (line 293) | def forward(self, x):
  class AutoencoderKL (line 304) | class AutoencoderKL(nn.Module):
    method __init__ (line 305) | def __init__(
    method dtype (line 357) | def dtype(self):
    method decode (line 360) | def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Au...

FILE: src/zimage/pipeline.py
  function calculate_shift (line 23) | def calculate_shift(
  function retrieve_timesteps (line 36) | def retrieve_timesteps(
  function generate (line 67) | def generate(

FILE: src/zimage/scheduler.py
  class SchedulerOutput (line 13) | class SchedulerOutput:
  class SchedulerConfig (line 17) | class SchedulerConfig:
    method __init__ (line 18) | def __init__(self, **kwargs):
    method get (line 21) | def get(self, key, default=None):
    method __getattr__ (line 24) | def __getattr__(self, name):
  class FlowMatchEulerDiscreteScheduler (line 28) | class FlowMatchEulerDiscreteScheduler:
    method __init__ (line 31) | def __init__(
    method set_timesteps (line 62) | def set_timesteps(
    method index_for_timestep (line 104) | def index_for_timestep(self, timestep, schedule_timesteps=None):
    method _init_step_index (line 112) | def _init_step_index(self, timestep):
    method step (line 120) | def step(
    method _sigma_to_t (line 146) | def _sigma_to_t(self, sigma):
    method time_shift (line 149) | def time_shift(self, mu: float, sigma: float, t: torch.Tensor):

FILE: src/zimage/transformer.py
  class TimestepEmbedder (line 22) | class TimestepEmbedder(nn.Module):
    method __init__ (line 23) | def __init__(self, out_size, mid_size=None, frequency_embedding_size=F...
    method timestep_embedding (line 35) | def timestep_embedding(t, dim, max_period=MAX_PERIOD):
    method forward (line 47) | def forward(self, t):
  class RMSNorm (line 56) | class RMSNorm(nn.Module):
    method __init__ (line 57) | def __init__(self, dim: int, eps: float = 1e-5):
    method forward (line 62) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class FeedForward (line 67) | class FeedForward(nn.Module):
    method __init__ (line 68) | def __init__(self, dim: int, hidden_dim: int):
    method forward (line 74) | def forward(self, x):
  function apply_rotary_emb (line 78) | def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> tor...
  class ZImageAttention (line 86) | class ZImageAttention(nn.Module):
    method __init__ (line 89) | def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: b...
    method forward (line 103) | def forward(
  class ZImageTransformerBlock (line 143) | class ZImageTransformerBlock(nn.Module):
    method __init__ (line 144) | def __init__(
    method forward (line 171) | def forward(
  class FinalLayer (line 205) | class FinalLayer(nn.Module):
    method __init__ (line 206) | def __init__(self, hidden_size, out_channels):
    method forward (line 215) | def forward(self, x, c):
  class RopeEmbedder (line 222) | class RopeEmbedder:
    method __init__ (line 223) | def __init__(
    method precompute_freqs_cis (line 236) | def precompute_freqs_cis(dim: List[int], end: List[int], theta: float ...
    method __call__ (line 247) | def __call__(self, ids: torch.Tensor):
  class ZImageTransformer2DModel (line 266) | class ZImageTransformer2DModel(nn.Module):
    method __init__ (line 267) | def __init__(
    method unpatchify (line 345) | def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_s...
    method create_coordinate_grid (line 362) | def create_coordinate_grid(size, start=None, device=None):
    method patchify_and_embed (line 369) | def patchify_and_embed(
    method forward (line 474) | def forward(