SYMBOL INDEX (813 symbols across 110 files)

FILE: autoregressive/models/dinov2_adapter.py
  class Dinov2_Adapter (line 8) | class Dinov2_Adapter(nn.Module):
    method __init__ (line 9) | def __init__(self, input_dim=1, output_dim=768, attention=False, pool=...
    method to_patch14 (line 16) | def to_patch14(self, input):
    method forward (line 26) | def forward(self, x):

FILE: autoregressive/models/generate.py
  function top_k_top_p_filtering (line 17) | def top_k_top_p_filtering(
  function sample (line 59) | def sample(logits, temperature: float=1.0, top_k: int=2000, top_p: float...
  function logits_to_probs (line 77) | def logits_to_probs(logits, temperature: float = 1.0, top_p: float=1.0, ...
  function prefill (line 85) | def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_...
  function decode_one_token (line 97) | def decode_one_token(model, x: torch.Tensor, input_pos: torch.Tensor, cf...
  function decode_n_tokens (line 113) | def decode_n_tokens(
  function generate (line 135) | def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0,...

FILE: autoregressive/models/gpt.py
  function get_causal_mask (line 19) | def get_causal_mask(seq_length):
  function find_multiple (line 25) | def find_multiple(n: int, k: int):
  class ModelArgs (line 31) | class ModelArgs:
  class LabelEmbedder (line 66) | class LabelEmbedder(nn.Module):
    method __init__ (line 70) | def __init__(self, num_classes, hidden_size, dropout_prob):
    method token_drop (line 77) | def token_drop(self, labels, force_drop_ids=None):
    method forward (line 88) | def forward(self, labels, train, force_drop_ids=None):
  class ConditionEmbedder (line 99) | class ConditionEmbedder(nn.Module):
    method __init__ (line 103) | def __init__(self, in_channels, hidden_size, uncond_prob, token_num=12...
    method token_drop (line 109) | def token_drop(self, caption, force_drop_ids=None, drop_ids=None):
    method forward (line 122) | def forward(self, caption, train, force_drop_ids=None, drop_ids=None):
  class CaptionEmbedder (line 132) | class CaptionEmbedder(nn.Module):
    method __init__ (line 136) | def __init__(self, in_channels, hidden_size, uncond_prob, token_num=120):
    method token_drop (line 142) | def token_drop(self, caption, force_drop_ids=None):
    method forward (line 153) | def forward(self, caption, train, force_drop_ids=None):
  class MLP (line 161) | class MLP(nn.Module):
    method __init__ (line 162) | def __init__(self, in_features, hidden_features, out_features):
    method forward (line 173) | def forward(self, x):
  class RMSNorm (line 183) | class RMSNorm(torch.nn.Module):
    method __init__ (line 184) | def __init__(self, dim: int, eps: float = 1e-5):
    method _norm (line 189) | def _norm(self, x):
    method forward (line 192) | def forward(self, x):
  class FeedForward (line 197) | class FeedForward(nn.Module):
    method __init__ (line 198) | def __init__(self, config: ModelArgs):
    method forward (line 212) | def forward(self, x):
  class KVCache (line 216) | class KVCache(nn.Module):
    method __init__ (line 217) | def __init__(self, max_batch_size, max_seq_length, n_head, head_dim, d...
    method update (line 223) | def update(self, input_pos, k_val, v_val):
  class Attention (line 234) | class Attention(nn.Module):
    method __init__ (line 235) | def __init__(self, config: ModelArgs):
    method forward (line 253) | def forward(
  class TransformerBlock (line 290) | class TransformerBlock(nn.Module):
    method __init__ (line 291) | def __init__(self, config: ModelArgs, drop_path: float):
    method forward (line 299) | def forward(
  class Transformer (line 306) | class Transformer(nn.Module):
    method __init__ (line 307) | def __init__(self, config: ModelArgs):
    method initialize_weights (line 368) | def initialize_weights(self):
    method _init_weights (line 373) | def _init_weights(self, module):
    method setup_caches (line 382) | def setup_caches(self, max_batch_size, max_seq_length, dtype):
    method forward (line 400) | def forward(
    method get_fsdp_wrap_module_list (line 468) | def get_fsdp_wrap_module_list(self) -> List[nn.Module]:
  function precompute_freqs_cis (line 477) | def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, c...
  function precompute_freqs_cis_2d (line 487) | def precompute_freqs_cis_2d(grid_size: int, n_elem: int, base: int = 100...
  function apply_rotary_emb (line 503) | def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
  function GPT_7B (line 521) | def GPT_7B(**kwargs):
  function GPT_3B (line 524) | def GPT_3B(**kwargs):
  function GPT_1B (line 527) | def GPT_1B(**kwargs):
  function GPT_XXXL (line 531) | def GPT_XXXL(**kwargs):
  function GPT_XXL (line 534) | def GPT_XXL(**kwargs):
  function GPT_XL (line 537) | def GPT_XL(**kwargs):
  function GPT_L (line 540) | def GPT_L(**kwargs):
  function GPT_B (line 543) | def GPT_B(**kwargs):

FILE: autoregressive/models/gpt_t2i.py
  function get_causal_mask (line 20) | def get_causal_mask(seq_length):
  function find_multiple (line 26) | def find_multiple(n: int, k: int):
  class ModelArgs (line 32) | class ModelArgs:
  class LabelEmbedder (line 67) | class LabelEmbedder(nn.Module):
    method __init__ (line 71) | def __init__(self, num_classes, hidden_size, dropout_prob):
    method token_drop (line 78) | def token_drop(self, labels, force_drop_ids=None):
    method forward (line 89) | def forward(self, labels, train, force_drop_ids=None):
  class ConditionEmbedder (line 100) | class ConditionEmbedder(nn.Module):
    method __init__ (line 104) | def __init__(self, in_channels, hidden_size, uncond_prob, token_num=12...
    method token_drop (line 110) | def token_drop(self, caption, force_drop_ids=None, drop_ids=None):
    method forward (line 123) | def forward(self, caption, train, force_drop_ids=None, drop_ids=None):
  class CaptionEmbedder (line 133) | class CaptionEmbedder(nn.Module):
    method __init__ (line 137) | def __init__(self, in_channels, hidden_size, uncond_prob, token_num=120):
    method token_drop (line 143) | def token_drop(self, caption, force_drop_ids=None):
    method forward (line 154) | def forward(self, caption, train, force_drop_ids=None):
  class MLP (line 165) | class MLP(nn.Module):
    method __init__ (line 166) | def __init__(self, in_features, hidden_features, out_features):
    method forward (line 177) | def forward(self, x):
  class RMSNorm (line 187) | class RMSNorm(torch.nn.Module):
    method __init__ (line 188) | def __init__(self, dim: int, eps: float = 1e-5):
    method _norm (line 193) | def _norm(self, x):
    method forward (line 196) | def forward(self, x):
  class FeedForward (line 201) | class FeedForward(nn.Module):
    method __init__ (line 202) | def __init__(self, config: ModelArgs):
    method forward (line 216) | def forward(self, x):
  class KVCache (line 220) | class KVCache(nn.Module):
    method __init__ (line 221) | def __init__(self, max_batch_size, max_seq_length, n_head, head_dim, d...
    method update (line 227) | def update(self, input_pos, k_val, v_val):
  class Attention (line 238) | class Attention(nn.Module):
    method __init__ (line 239) | def __init__(self, config: ModelArgs):
    method forward (line 257) | def forward(
  class TransformerBlock (line 294) | class TransformerBlock(nn.Module):
    method __init__ (line 295) | def __init__(self, config: ModelArgs, drop_path: float):
    method forward (line 303) | def forward(
  class Transformer (line 310) | class Transformer(nn.Module):
    method __init__ (line 311) | def __init__(self, config: ModelArgs):
    method initialize_weights (line 372) | def initialize_weights(self):
    method _init_weights (line 381) | def _init_weights(self, module):
    method setup_caches (line 391) | def setup_caches(self, max_batch_size, max_seq_length, dtype):
    method forward (line 409) | def forward(
    method get_fsdp_wrap_module_list (line 487) | def get_fsdp_wrap_module_list(self) -> List[nn.Module]:
  function precompute_freqs_cis (line 496) | def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, c...
  function precompute_freqs_cis_2d (line 506) | def precompute_freqs_cis_2d(grid_size: int, n_elem: int, base: int = 100...
  function apply_rotary_emb (line 522) | def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
  function GPT_7B (line 540) | def GPT_7B(**kwargs):
  function GPT_3B (line 543) | def GPT_3B(**kwargs):
  function GPT_1B (line 546) | def GPT_1B(**kwargs):
  function GPT_XXXL (line 550) | def GPT_XXXL(**kwargs):
  function GPT_XXL (line 553) | def GPT_XXL(**kwargs):
  function GPT_XL (line 556) | def GPT_XL(**kwargs):
  function GPT_L (line 559) | def GPT_L(**kwargs):
  function GPT_B (line 562) | def GPT_B(**kwargs):

FILE: autoregressive/models/vit_adapter.py
  class ViT_Adapter (line 8) | class ViT_Adapter(nn.Module):
    method __init__ (line 9) | def __init__(self, input_dim=3, output_dim=768, attention=False, pool=...
    method forward (line 13) | def forward(self, x):

FILE: autoregressive/sample/sample_c2i.py
  function main (line 27) | def main(args):

FILE: autoregressive/sample/sample_c2i_ddp.py
  function create_npz_from_sample_folder (line 21) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function main (line 38) | def main(args):

FILE: autoregressive/sample/sample_t2i.py
  function main (line 34) | def main(args):

FILE: autoregressive/sample/sample_t2i_MR.py
  function resize_image_to_16_multiple (line 37) | def resize_image_to_16_multiple(image_path, condition_type='seg'):
  function main (line 51) | def main(args):

FILE: autoregressive/sample/sample_t2i_ddp.py
  function main (line 26) | def main(args):

FILE: autoregressive/serve/gpt_model.py
  function find_multiple (line 18) | def find_multiple(n: int, k: int):
  class ModelArgs (line 24) | class ModelArgs:
  class LabelEmbedder (line 50) | class LabelEmbedder(nn.Module):
    method __init__ (line 54) | def __init__(self, num_classes, hidden_size, dropout_prob):
    method forward (line 73) | def forward(self, labels):
  class FeedForward (line 98) | class FeedForward(nn.Module):
    method __init__ (line 99) | def __init__(self, config: ModelArgs):
    method forward (line 119) | def forward(self, x):
  class Attention (line 127) | class Attention(nn.Module):
    method __init__ (line 128) | def __init__(self, config: ModelArgs):
    method forward (line 154) | def forward(
  class AttentionMonkeyPatch (line 178) | class AttentionMonkeyPatch(Attention):
    method __init__ (line 185) | def __init__(self, config: ModelArgs):
    method forward (line 191) | def forward(
  class TransformerBlock (line 227) | class TransformerBlock(nn.Module):
    method __init__ (line 228) | def __init__(self, config: ModelArgs):
    method forward (line 238) | def forward(self, x: torch.Tensor, positions: torch.Tensor, kv_cache: ...
  class Transformer (line 244) | class Transformer(nn.Module):
    method __init__ (line 245) | def __init__(self, config: ModelArgs):
    method forward (line 273) | def forward(
    method compute_logits (line 297) | def compute_logits(self, hidden_states: torch.Tensor,
    method sample (line 302) | def sample(
    method custom_load_state_dict (line 311) | def custom_load_state_dict(self, model_weights):
  function precompute_freqs_cis (line 334) | def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, c...
  function precompute_freqs_cis_2d (line 344) | def precompute_freqs_cis_2d(grid_size: int, n_elem: int, base: int = 100...
  function apply_rotary_emb (line 360) | def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
  function apply_rotary_emb_bs (line 373) | def apply_rotary_emb_bs(x: torch.Tensor, freqs_cis: torch.Tensor):
  function GPT_7B (line 390) | def GPT_7B(**kwargs):
  function GPT_3B (line 393) | def GPT_3B(**kwargs):
  function GPT_1B (line 396) | def GPT_1B(**kwargs):
  function GPT_XXXL (line 400) | def GPT_XXXL(**kwargs):
  function GPT_XXL (line 403) | def GPT_XXL(**kwargs):
  function GPT_XL (line 406) | def GPT_XL(**kwargs):
  function GPT_L (line 409) | def GPT_L(**kwargs):
  function GPT_B (line 412) | def GPT_B(**kwargs):

FILE: autoregressive/serve/gpu_executor.py
  class GPUExecutor (line 17) | class GPUExecutor(ExecutorBase):
    method __init__ (line 18) | def __init__(
    method _init_executor (line 44) | def _init_executor(self) -> None:
    method _init_non_spec_worker (line 55) | def _init_non_spec_worker(self):
    method _init_spec_worker (line 83) | def _init_spec_worker(self):
    method determine_num_available_blocks (line 136) | def determine_num_available_blocks(self) -> Tuple[int, int]:
    method initialize_cache (line 142) | def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:
    method execute_model (line 153) | def execute_model(
    method add_lora (line 170) | def add_lora(self, lora_request: LoRARequest) -> bool:
    method remove_lora (line 174) | def remove_lora(self, lora_id: int) -> bool:
    method list_loras (line 178) | def list_loras(self) -> Set[int]:
    method check_health (line 181) | def check_health(self) -> None:
  class GPUExecutorAsync (line 187) | class GPUExecutorAsync(GPUExecutor, ExecutorAsyncBase):
    method execute_model_async (line 189) | async def execute_model_async(

FILE: autoregressive/serve/llm.py
  class LLM (line 22) | class LLM:
    method __init__ (line 82) | def __init__(
    method get_tokenizer (line 128) | def get_tokenizer(
    method set_tokenizer (line 132) | def set_tokenizer(
    method generate (line 138) | def generate(
    method _add_request (line 221) | def _add_request(
    method _run_engine (line 238) | def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]:

FILE: autoregressive/serve/llm_engine.py
  function _load_generation_config_dict (line 40) | def _load_generation_config_dict(model_config: ModelConfig):
  class LLMEngine (line 51) | class LLMEngine:
    method __init__ (line 85) | def __init__(
    method _initialize_kv_caches (line 234) | def _initialize_kv_caches(self) -> None:
    method from_engine_args (line 255) | def from_engine_args(
    method __reduce__ (line 293) | def __reduce__(self):
    method get_tokenizer (line 298) | def get_tokenizer(self) -> "PreTrainedTokenizer":
    method get_tokenizer_for_seq (line 301) | def get_tokenizer_for_seq(self,
    method _init_tokenizer (line 305) | def _init_tokenizer(self, **tokenizer_init_kwargs):
    method _verify_args (line 318) | def _verify_args(self) -> None:
    method encode_request (line 326) | def encode_request(
    method add_request (line 340) | def add_request(
    method abort_request (line 439) | def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
    method get_model_config (line 458) | def get_model_config(self) -> ModelConfig:
    method get_num_unfinished_requests (line 462) | def get_num_unfinished_requests(self) -> int:
    method has_unfinished_requests (line 466) | def has_unfinished_requests(self) -> bool:
    method _process_model_outputs (line 470) | def _process_model_outputs(
    method step (line 511) | def step(self) -> List[RequestOutput]:
    method do_log_stats (line 583) | def do_log_stats(self) -> None:
    method _get_stats (line 588) | def _get_stats(self,
    method add_lora (line 661) | def add_lora(self, lora_request: LoRARequest) -> bool:
    method remove_lora (line 664) | def remove_lora(self, lora_id: int) -> bool:
    method list_loras (line 667) | def list_loras(self) -> List[int]:
    method check_health (line 670) | def check_health(self) -> None:

FILE: autoregressive/serve/model_runner.py
  class PreparePromptMetadata (line 43) | class PreparePromptMetadata(NamedTuple):
    method empty (line 56) | def empty(cls):
  class PrepareDecodeMetadata (line 71) | class PrepareDecodeMetadata(NamedTuple):
    method empty (line 81) | def empty(cls):
  class BatchType (line 94) | class BatchType(IntEnum):
  class ModelRunner (line 103) | class ModelRunner:
    method __init__ (line 105) | def __init__(
    method load_model (line 161) | def load_model(self, args) -> None:
    method set_block_size (line 237) | def set_block_size(self, block_size: int) -> None:
    method get_max_block_per_batch (line 244) | def get_max_block_per_batch(self) -> int:
    method _prepare_prompt (line 248) | def _prepare_prompt(
    method _prepare_decode (line 448) | def _prepare_decode(
    method _prepare_sample (line 574) | def _prepare_sample(
    method prepare_input_tensors (line 676) | def prepare_input_tensors(
    method execute_model (line 845) | def execute_model(
    method profile_run (line 889) | def profile_run(self) -> None:
    method remove_all_loras (line 955) | def remove_all_loras(self) -> bool:
    method set_active_loras (line 960) | def set_active_loras(self, lora_requests: Set[LoRARequest],
    method add_lora (line 966) | def add_lora(self, lora_request: LoRARequest) -> bool:
    method remove_lora (line 971) | def remove_lora(self, lora_id: int) -> bool:
    method list_loras (line 976) | def list_loras(self) -> Set[int]:
    method capture_model (line 982) | def capture_model(self, kv_caches: List[torch.Tensor]) -> None:
    method __del__ (line 1083) | def __del__(self) -> None:
    method vocab_size (line 1094) | def vocab_size(self) -> int:
  class CUDAGraphRunner (line 1098) | class CUDAGraphRunner:
    method __init__ (line 1100) | def __init__(self, model: nn.Module):
    method graph (line 1108) | def graph(self):
    method capture (line 1112) | def capture(
    method forward (line 1162) | def forward(
    method __call__ (line 1188) | def __call__(self, *args, **kwargs):
  function _maybe_pynccl (line 1193) | def _maybe_pynccl():
  function _get_graph_batch_size (line 1202) | def _get_graph_batch_size(batch_size: int) -> int:
  function _prepare_fake_inputs (line 1217) | def _prepare_fake_inputs(

FILE: autoregressive/serve/sample_c2i.py
  function main (line 13) | def main(args):

FILE: autoregressive/serve/sampler.py
  class Sampler (line 17) | class Sampler(nn.Module):
    method __init__ (line 38) | def __init__(self, cfg_scale=1.0):
    method forward (line 46) | def forward(
    method _should_modify_greedy_probs_inplace (line 128) | def _should_modify_greedy_probs_inplace(self) -> bool:
  function _get_bin_counts_and_mask (line 143) | def _get_bin_counts_and_mask(
  function _apply_min_tokens_penalty (line 160) | def _apply_min_tokens_penalty(
  function _apply_penalties (line 207) | def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.T...
  function _apply_top_k_top_p (line 230) | def _apply_top_k_top_p(
  function _apply_min_p (line 262) | def _apply_min_p(
  function _greedy_sample (line 279) | def _greedy_sample(
  function _random_sample (line 298) | def _random_sample(
  function _beam_search_sample (line 325) | def _beam_search_sample(
  function _multinomial (line 383) | def _multinomial(
  function _sample_with_torch (line 410) | def _sample_with_torch(
  function _sample_with_triton_kernel (line 520) | def _sample_with_triton_kernel(
  function _sample (line 600) | def _sample(
  function _get_ranks (line 618) | def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
  function _get_logprobs (line 637) | def _get_logprobs(
  function _modify_greedy_probs_inplace (line 775) | def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Te...
  function _build_sampler_output (line 826) | def _build_sampler_output(

FILE: autoregressive/serve/worker.py
  class Worker (line 27) | class Worker(WorkerBase):
    method __init__ (line 35) | def __init__(
    method init_device (line 89) | def init_device(self) -> None:
    method load_model (line 117) | def load_model(self, args):
    method determine_num_available_blocks (line 121) | def determine_num_available_blocks(self) -> Tuple[int, int]:
    method initialize_cache (line 166) | def initialize_cache(self, num_gpu_blocks: int,
    method _init_cache_engine (line 182) | def _init_cache_engine(self):
    method _warm_up_model (line 189) | def _warm_up_model(self) -> None:
    method cache_swap (line 196) | def cache_swap(
    method execute_model (line 212) | def execute_model(
    method add_lora (line 257) | def add_lora(self, lora_request: LoRARequest) -> bool:
    method remove_lora (line 260) | def remove_lora(self, lora_id: int) -> bool:
    method list_loras (line 263) | def list_loras(self) -> Set[int]:
    method max_model_len (line 267) | def max_model_len(self) -> int:
    method vocab_size (line 271) | def vocab_size(self) -> int:
    method get_cache_block_size_bytes (line 274) | def get_cache_block_size_bytes(self) -> int:
  function init_worker_distributed_environment (line 282) | def init_worker_distributed_environment(
  function _check_if_gpu_supports_dtype (line 322) | def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
  function raise_if_cache_size_invalid (line 336) | def raise_if_cache_size_invalid(num_gpu_blocks, block_size,

FILE: autoregressive/test/metric.py
  class SSIM (line 7) | class SSIM:
    method __init__ (line 8) | def __init__(self, data_range=1.0):
    method update (line 12) | def update(self, img1, img2):
    method calculate (line 17) | def calculate(self):
  class F1score (line 24) | class F1score:
    method __init__ (line 25) | def __init__(self, threshold=128):
    method update (line 30) | def update(self, img1, img2):
    method calculate (line 45) | def calculate(self):
  class RMSE (line 49) | class RMSE:
    method __init__ (line 50) | def __init__(self):
    method update (line 54) | def update(self, img1, img2):
    method calculate (line 64) | def calculate(self):

FILE: autoregressive/test/test_c2i.py
  function create_npz_from_sample_folder (line 37) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function main (line 53) | def main(args):

FILE: autoregressive/test/test_t2i.py
  function main (line 43) | def main(args):

FILE: autoregressive/train/extract_codes_c2i.py
  function main (line 25) | def main(args):

FILE: autoregressive/train/extract_codes_t2i.py
  class CustomDataset (line 25) | class CustomDataset(Dataset):
    method __init__ (line 26) | def __init__(self, lst_dir, start, end, transform):
    method __len__ (line 41) | def __len__(self):
    method __getitem__ (line 44) | def __getitem__(self, index):
  function main (line 56) | def main(args):

FILE: autoregressive/train/extract_file_ade.py
  function collate_fn (line 53) | def collate_fn(examples):
  function main (line 79) | def main(args):

FILE: autoregressive/train/extract_file_cocostuff.py
  function collate_fn (line 53) | def collate_fn(examples):
  function main (line 79) | def main(args):

FILE: autoregressive/train/extract_file_imagenet.py
  function main (line 29) | def main(args):

FILE: autoregressive/train/extract_file_multigen.py
  function collate_fn (line 59) | def collate_fn(examples):
  function main (line 79) | def main(args):

FILE: autoregressive/train/train_c2i.py
  function creat_optimizer (line 28) | def creat_optimizer(model, weight_decay, learning_rate, betas, logger):
  function main (line 57) | def main(args):

FILE: autoregressive/train/train_c2i_canny.py
  function creat_optimizer (line 31) | def creat_optimizer(model, weight_decay, learning_rate, betas, logger):
  function main (line 59) | def main(args):

FILE: autoregressive/train/train_c2i_depth.py
  function creat_optimizer (line 33) | def creat_optimizer(model, weight_decay, learning_rate, betas, logger):
  function main (line 60) | def main(args):

FILE: autoregressive/train/train_c2i_fsdp.py
  function setup_fsdp_sync (line 31) | def setup_fsdp_sync(model: nn.Module, args: argparse.Namespace, device) ...
  function creat_optimizer_by_name (line 67) | def creat_optimizer_by_name(model, weight_decay, learning_rate, betas, g...
  function main (line 102) | def main(args):

FILE: autoregressive/train/train_t2i.py
  function main (line 27) | def main(args):

FILE: autoregressive/train/train_t2i_canny.py
  function main (line 38) | def main(args):

FILE: autoregressive/train/train_t2i_depth.py
  function main (line 38) | def main(args):

FILE: autoregressive/train/train_t2i_depth_multiscale.py
  function random_sample_scale (line 44) | def random_sample_scale(image, condition=None):
  function main (line 59) | def main(args):

FILE: autoregressive/train/train_t2i_hed.py
  function main (line 38) | def main(args):

FILE: autoregressive/train/train_t2i_hed_multiscale.py
  function random_sample_scale (line 42) | def random_sample_scale(image, condition=None):
  function main (line 57) | def main(args):

FILE: autoregressive/train/train_t2i_lineart.py
  function main (line 39) | def main(args):

FILE: autoregressive/train/train_t2i_lineart_multiscale.py
  function random_sample_scale (line 44) | def random_sample_scale(image, condition=None):
  function main (line 58) | def main(args):

FILE: autoregressive/train/train_t2i_seg.py
  function main (line 39) | def main(args):

FILE: autoregressive/train/train_t2i_seg_multiscale.py
  function random_sample_scale (line 43) | def random_sample_scale(image, condition=None):
  function main (line 57) | def main(args):

FILE: condition/canny.py
  class CannyDetector (line 6) | class CannyDetector:
    method __call__ (line 7) | def __call__(self, img, low_threshold=100, high_threshold=200):

FILE: condition/depth.py
  class Depth (line 6) | class Depth:
    method __init__ (line 7) | def __init__(self, device):
    method __call__ (line 10) | def __call__(self, input_image):

FILE: condition/hed.py
  class DoubleConvBlock (line 17) | class DoubleConvBlock(torch.nn.Module):
    method __init__ (line 18) | def __init__(self, input_channel, output_channel, layer_number):
    method __call__ (line 26) | def __call__(self, x, down_sampling=False):
  class ControlNetHED_Apache2 (line 36) | class ControlNetHED_Apache2(torch.nn.Module):
    method __init__ (line 37) | def __init__(self):
    method __call__ (line 46) | def __call__(self, x):
  class HEDdetector (line 56) | class HEDdetector(torch.nn.Module):
    method __init__ (line 57) | def __init__(self):
    method __call__ (line 67) | def __call__(self, input_image):
  function nms (line 84) | def nms(x, t, s):

FILE: condition/lineart.py
  class ResidualBlock (line 9) | class ResidualBlock(nn.Module):
    method __init__ (line 10) | def __init__(self, in_features):
    method forward (line 24) | def forward(self, x):
  class LineArt (line 26) | class LineArt(nn.Module):
    method __init__ (line 27) | def __init__(self, input_nc=3, output_nc=1, n_residual_blocks=3, sigmo...
    method forward (line 74) | def forward(self, x, cond=None):

FILE: condition/midas/depth.py
  function disabled_train (line 31) | def disabled_train(self, mode=True):
  function load_midas_transform (line 37) | def load_midas_transform(model_type):
  function load_model (line 82) | def load_model(model_type):
  class MiDaSInference (line 150) | class MiDaSInference(nn.Module):
    method __init__ (line 163) | def __init__(self, model_type):
    method forward (line 170) | def forward(self, x):
  class MidasDetector (line 176) | class MidasDetector:
    method __init__ (line 177) | def __init__(self,device=torch.device('cuda:0'), model_type="dpt_hybri...
    method __call__ (line 181) | def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):

FILE: condition/midas/midas/base_model.py
  class BaseModel (line 4) | class BaseModel(torch.nn.Module):
    method load (line 5) | def load(self, path):

FILE: condition/midas/midas/blocks.py
  function _make_encoder (line 11) | def _make_encoder(backbone, features, use_pretrained, groups=1, expand=F...
  function _make_scratch (line 49) | def _make_scratch(in_shape, out_shape, groups=1, expand=False):
  function _make_pretrained_efficientnet_lite3 (line 78) | def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
  function _make_efficientnet_backbone (line 88) | def _make_efficientnet_backbone(effnet):
  function _make_resnet_backbone (line 101) | def _make_resnet_backbone(resnet):
  function _make_pretrained_resnext101_wsl (line 114) | def _make_pretrained_resnext101_wsl(use_pretrained):
  class Interpolate (line 120) | class Interpolate(nn.Module):
    method __init__ (line 124) | def __init__(self, scale_factor, mode, align_corners=False):
    method forward (line 138) | def forward(self, x):
  class ResidualConvUnit (line 155) | class ResidualConvUnit(nn.Module):
    method __init__ (line 159) | def __init__(self, features):
    method forward (line 177) | def forward(self, x):
  class FeatureFusionBlock (line 194) | class FeatureFusionBlock(nn.Module):
    method __init__ (line 198) | def __init__(self, features):
    method forward (line 209) | def forward(self, *xs):
  class ResidualConvUnit_custom (line 231) | class ResidualConvUnit_custom(nn.Module):
    method __init__ (line 235) | def __init__(self, features, activation, bn):
    method forward (line 263) | def forward(self, x):
  class FeatureFusionBlock_custom (line 291) | class FeatureFusionBlock_custom(nn.Module):
    method __init__ (line 295) | def __init__(self, features, activation, deconv=False, bn=False, expan...
    method forward (line 320) | def forward(self, *xs):

FILE: condition/midas/midas/dpt_depth.py
  function _make_fusion_block (line 15) | def _make_fusion_block(features, use_bn):
  class DPT (line 26) | class DPT(BaseModel):
    method __init__ (line 27) | def __init__(
    method forward (line 67) | def forward(self, x):
  class DPTDepthModel (line 88) | class DPTDepthModel(DPT):
    method __init__ (line 89) | def __init__(self, path=None, non_negative=True, **kwargs):
    method forward (line 107) | def forward(self, x):

FILE: condition/midas/midas/midas_net.py
  class MidasNet (line 12) | class MidasNet(BaseModel):
    method __init__ (line 16) | def __init__(self, path=None, features=256, non_negative=True):
    method forward (line 49) | def forward(self, x):

FILE: condition/midas/midas/midas_net_custom.py
  class MidasNet_small (line 12) | class MidasNet_small(BaseModel):
    method __init__ (line 16) | def __init__(self, path=None, features=64, backbone="efficientnet_lite...
    method forward (line 73) | def forward(self, x):
  function fuse_model (line 109) | def fuse_model(m):

FILE: condition/midas/midas/transforms.py
  function apply_min_size (line 6) | def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AR...
  class Resize (line 48) | class Resize(object):
    method __init__ (line 52) | def __init__(
    method constrain_to_multiple_of (line 94) | def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
    method get_size (line 105) | def get_size(self, width, height):
    method __call__ (line 162) | def __call__(self, sample):
  class NormalizeImage (line 197) | class NormalizeImage(object):
    method __init__ (line 201) | def __init__(self, mean, std):
    method __call__ (line 205) | def __call__(self, sample):
  class PrepareForNet (line 211) | class PrepareForNet(object):
    method __init__ (line 215) | def __init__(self):
    method __call__ (line 218) | def __call__(self, sample):

FILE: condition/midas/midas/vit.py
  class Slice (line 9) | class Slice(nn.Module):
    method __init__ (line 10) | def __init__(self, start_index=1):
    method forward (line 14) | def forward(self, x):
  class AddReadout (line 18) | class AddReadout(nn.Module):
    method __init__ (line 19) | def __init__(self, start_index=1):
    method forward (line 23) | def forward(self, x):
  class ProjectReadout (line 31) | class ProjectReadout(nn.Module):
    method __init__ (line 32) | def __init__(self, in_features, start_index=1):
    method forward (line 38) | def forward(self, x):
  class Transpose (line 45) | class Transpose(nn.Module):
    method __init__ (line 46) | def __init__(self, dim0, dim1):
    method forward (line 51) | def forward(self, x):
  function forward_vit (line 56) | def forward_vit(pretrained, x):
  function _resize_pos_embed (line 100) | def _resize_pos_embed(self, posemb, gs_h, gs_w):
  function forward_flex (line 117) | def forward_flex(self, x):
  function get_activation (line 159) | def get_activation(name):
  function get_readout_oper (line 166) | def get_readout_oper(vit_features, features, use_readout, start_index=1):
  function _make_vit_b16_backbone (line 183) | def _make_vit_b16_backbone(
  function _make_pretrained_vitl16_384 (line 297) | def _make_pretrained_vitl16_384(pretrained, use_readout="ignore", hooks=...
  function _make_pretrained_vitb16_384 (line 310) | def _make_pretrained_vitb16_384(pretrained, use_readout="ignore", hooks=...
  function _make_pretrained_deitb16_384 (line 319) | def _make_pretrained_deitb16_384(pretrained, use_readout="ignore", hooks...
  function _make_pretrained_deitb16_distil_384 (line 328) | def _make_pretrained_deitb16_distil_384(pretrained, use_readout="ignore"...
  function _make_vit_b_rn50_backbone (line 343) | def _make_vit_b_rn50_backbone(
  function _make_pretrained_vitb_rn50_384 (line 478) | def _make_pretrained_vitb_rn50_384(

FILE: condition/utils.py
  function HWC3 (line 9) | def HWC3(x):
  function resize_image (line 28) | def resize_image(input_image, resolution):

FILE: create_npz.py
  function create_npz_from_sample_folder (line 8) | def create_npz_from_sample_folder(sample_dir, num=50_000):

FILE: dataset/augmentation.py
  function center_crop_arr (line 8) | def center_crop_arr(pil_image, image_size):
  function random_crop_arr (line 29) | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_f...

FILE: dataset/build.py
  function build_dataset (line 8) | def build_dataset(args, **kwargs):

FILE: dataset/coco.py
  class SingleFolderDataset (line 7) | class SingleFolderDataset(Dataset):
    method __init__ (line 8) | def __init__(self, directory, transform=None):
    method __len__ (line 15) | def __len__(self):
    method __getitem__ (line 18) | def __getitem__(self, idx):
  function build_coco (line 26) | def build_coco(args, transform):

FILE: dataset/imagenet.py
  class CustomDataset (line 9) | class CustomDataset(Dataset):
    method __init__ (line 10) | def __init__(self, feature_dir, label_dir, condition_dir=None, get_con...
    method __len__ (line 47) | def __len__(self):
    method __getitem__ (line 52) | def __getitem__(self, idx):
  function build_imagenet (line 104) | def build_imagenet(args, transform):
  function build_imagenet_code (line 107) | def build_imagenet_code(args):

FILE: dataset/openimage.py
  class DatasetJson (line 10) | class DatasetJson(Dataset):
    method __init__ (line 11) | def __init__(self, data_path, transform=None):
    method __len__ (line 20) | def __len__(self):
    method __getitem__ (line 23) | def __getitem__(self, idx):
    method getdata (line 32) | def getdata(self, idx):
  function build_openimage (line 41) | def build_openimage(args, transform):

FILE: dataset/pexels.py
  function build_pexels (line 3) | def build_pexels(args, transform):

FILE: dataset/t2i.py
  class Text2ImgDatasetImg (line 10) | class Text2ImgDatasetImg(Dataset):
    method __init__ (line 11) | def __init__(self, lst_dir, face_lst_dir, transform):
    method __len__ (line 39) | def __len__(self):
    method __getitem__ (line 42) | def __getitem__(self, index):
  class Text2ImgDataset (line 50) | class Text2ImgDataset(Dataset):
    method __init__ (line 51) | def __init__(self, args, transform):
    method __len__ (line 85) | def __len__(self):
    method dummy_data (line 88) | def dummy_data(self):
    method __getitem__ (line 95) | def __getitem__(self, index):
  class Text2ImgDatasetCode (line 138) | class Text2ImgDatasetCode(Dataset):
    method __init__ (line 139) | def __init__(self, args):
  function build_t2i_image (line 145) | def build_t2i_image(args, transform):
  function build_t2i (line 148) | def build_t2i(args, transform):
  function build_t2i_code (line 151) | def build_t2i_code(args):

FILE: dataset/t2i_control.py
  class T2IControlCode (line 36) | class T2IControlCode(Dataset):
    method __init__ (line 37) | def __init__(self, args):
    method __len__ (line 64) | def __len__(self):
    method dummy_data (line 67) | def dummy_data(self):
    method collate_fn (line 74) | def collate_fn(self, examples):
    method __getitem__ (line 104) | def __getitem__(self, index):
  function build_t2i_control_code (line 165) | def build_t2i_control_code(args):

FILE: dataset/utils.py
  function get_reward_model (line 20) | def get_reward_model(task='segmentation', model_path='mmseg::upernet/upe...
  function get_reward_loss (line 44) | def get_reward_loss(predictions, labels, task='segmentation', **args):
  function image_grid (line 65) | def image_grid(imgs, rows, cols):
  function map_color_to_index (line 77) | def map_color_to_index(image, dataset='limingcv/Captioned_ADE20K'):
  function seg_label_transform (line 105) | def seg_label_transform(
  function depth_label_transform (line 142) | def depth_label_transform(
  function edge_label_transform (line 156) | def edge_label_transform(labels, dataset_name):
  function label_transform (line 160) | def label_transform(labels, task, dataset_name, **args):
  function group_random_crop (line 171) | def group_random_crop(images, resolution):
  class ResidualBlock (line 191) | class ResidualBlock(nn.Module):
    method __init__ (line 192) | def __init__(self, in_features):
    method forward (line 206) | def forward(self, x):
  class LineDrawingModel (line 210) | class LineDrawingModel(nn.Module):
    method __init__ (line 211) | def __init__(self, input_nc=3, output_nc=1, n_residual_blocks=3, sigmo...
    method forward (line 258) | def forward(self, x, cond=None):
  class DoubleConvBlock (line 269) | class DoubleConvBlock(torch.nn.Module):
    method __init__ (line 270) | def __init__(self, input_channel, output_channel, layer_number):
    method __call__ (line 278) | def __call__(self, x, down_sampling=False):
  class ControlNetHED_Apache2 (line 288) | class ControlNetHED_Apache2(torch.nn.Module):
    method __init__ (line 289) | def __init__(self):
    method __call__ (line 298) | def __call__(self, x):
  class HEDdetector (line 308) | class HEDdetector(nn.Module):
    method __init__ (line 309) | def __init__(self, model_path):
    method __call__ (line 316) | def __call__(self, input_image):

FILE: demo/app_depth.py
  function randomize_seed_fn (line 5) | def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
  function create_demo (line 27) | def create_demo(process):

FILE: demo/app_edge.py
  function randomize_seed_fn (line 5) | def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
  function create_demo (line 27) | def create_demo(process):

FILE: demo/model.py
  class Model (line 21) | class Model:
    method __init__ (line 22) | def __init__(self):
    method to (line 34) | def to(self, device):
    method load_vq (line 37) | def load_vq(self):
    method load_gpt (line 48) | def load_gpt(self, condition_type='edge'):
    method load_gpt_weight (line 66) | def load_gpt_weight(self, condition_type='edge'):
    method load_t5 (line 77) | def load_t5(self):
    method process_edge (line 92) | def process_edge(
    method process_depth (line 192) | def process_depth(

FILE: evaluations/ade20k_mIoU.py
  function main (line 9) | def main():

FILE: evaluations/c2i/evaluator.py
  function main (line 27) | def main():
  class InvalidFIDException (line 75) | class InvalidFIDException(Exception):
  class FIDStatistics (line 79) | class FIDStatistics:
    method __init__ (line 80) | def __init__(self, mu: np.ndarray, sigma: np.ndarray):
    method frechet_distance (line 84) | def frechet_distance(self, other, eps=1e-6):
  class Evaluator (line 130) | class Evaluator:
    method __init__ (line 131) | def __init__(
    method warmup (line 147) | def warmup(self):
    method read_activations (line 150) | def read_activations(self, npz_path: str) -> Tuple[np.ndarray, np.ndar...
    method compute_activations (line 154) | def compute_activations(self, batches: Iterable[np.ndarray]) -> Tuple[...
    method read_statistics (line 176) | def read_statistics(
    method compute_statistics (line 186) | def compute_statistics(self, activations: np.ndarray) -> FIDStatistics:
    method compute_inception_score (line 191) | def compute_inception_score(self, activations: np.ndarray, split_size:...
    method compute_prec_recall (line 206) | def compute_prec_recall(
  class ManifoldEstimator (line 217) | class ManifoldEstimator:
    method __init__ (line 224) | def __init__(
    method warmup (line 253) | def warmup(self):
    method manifold_radii (line 260) | def manifold_radii(self, features: np.ndarray) -> np.ndarray:
    method evaluate (line 295) | def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_featu...
    method evaluate_pr (line 337) | def evaluate_pr(
  class DistanceBlock (line 374) | class DistanceBlock:
    method __init__ (line 381) | def __init__(self, session):
    method pairwise_distances (line 405) | def pairwise_distances(self, U, V):
    method less_thans (line 414) | def less_thans(self, batch_1, radii_1, batch_2, radii_2):
  function _batch_pairwise_distances (line 426) | def _batch_pairwise_distances(U, V):
  class NpzArrayReader (line 445) | class NpzArrayReader(ABC):
    method read_batch (line 447) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
    method remaining (line 451) | def remaining(self) -> int:
    method read_batches (line 454) | def read_batches(self, batch_size: int) -> Iterable[np.ndarray]:
  class BatchIterator (line 467) | class BatchIterator:
    method __init__ (line 468) | def __init__(self, gen_fn, length):
    method __len__ (line 472) | def __len__(self):
    method __iter__ (line 475) | def __iter__(self):
  class StreamingNpzArrayReader (line 479) | class StreamingNpzArrayReader(NpzArrayReader):
    method __init__ (line 480) | def __init__(self, arr_f, shape, dtype):
    method read_batch (line 486) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
    method remaining (line 501) | def remaining(self) -> int:
  class MemoryNpzArrayReader (line 505) | class MemoryNpzArrayReader(NpzArrayReader):
    method __init__ (line 506) | def __init__(self, arr):
    method load (line 511) | def load(cls, path: str, arr_name: str):
    method read_batch (line 516) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
    method remaining (line 524) | def remaining(self) -> int:
  function open_npz_array (line 529) | def open_npz_array(path: str, arr_name: str) -> NpzArrayReader:
  function _read_bytes (line 546) | def _read_bytes(fp, size, error_template="ran out of data"):
  function _open_npy_file (line 576) | def _open_npy_file(path: str, arr_name: str):
  function _download_inception_model (line 585) | def _download_inception_model():
  function _create_feature_graph (line 598) | def _create_feature_graph(input_batch):
  function _create_softmax_graph (line 615) | def _create_softmax_graph(input_batch):
  function _update_shapes (line 629) | def _update_shapes(pool3):
  function _numpy_partition (line 648) | def _numpy_partition(arr, kth, **kwargs):

FILE: evaluations/canny_f1score.py
  class ImageDataset (line 18) | class ImageDataset(Dataset):
    method __init__ (line 19) | def __init__(self, img_dir, label_dir):
    method __len__ (line 24) | def __len__(self):
    method __getitem__ (line 27) | def __getitem__(self, idx):

FILE: evaluations/clean_fid.py
  function main (line 4) | def main(args):

FILE: evaluations/cocostuff_mIoU.py
  function main (line 9) | def main():

FILE: evaluations/depth_rmse.py
  class ImageDataset (line 15) | class ImageDataset(Dataset):
    method __init__ (line 16) | def __init__(self, img_dir, label_dir):
    method __len__ (line 21) | def __len__(self):
    method __getitem__ (line 24) | def __getitem__(self, idx):

FILE: evaluations/hed_ssim.py
  class ImageDataset (line 17) | class ImageDataset(Dataset):
    method __init__ (line 18) | def __init__(self, img_dir, label_dir):
    method __len__ (line 23) | def __len__(self):
    method __getitem__ (line 26) | def __getitem__(self, idx):

FILE: evaluations/lineart_ssim.py
  class ImageDataset (line 18) | class ImageDataset(Dataset):
    method __init__ (line 19) | def __init__(self, img_dir, label_dir):
    method __len__ (line 24) | def __len__(self):
    method __getitem__ (line 27) | def __getitem__(self, idx):

FILE: evaluations/t2i/evaluation.py
  class CenterCropLongEdge (line 31) | class CenterCropLongEdge(object):
    method __call__ (line 37) | def __call__(self, img):
    method __repr__ (line 40) | def __repr__(self):
  class EvalDataset (line 44) | class EvalDataset(Dataset):
    method __init__ (line 45) | def __init__(self,
    method natural_sort (line 77) | def natural_sort(self, l):
    method load_dataset (line 82) | def load_dataset(self):
    method __len__ (line 100) | def __len__(self):
    method __getitem__ (line 104) | def __getitem__(self, index):
  function tensor2pil (line 117) | def tensor2pil(image: torch.Tensor):
  function compute_clip_score (line 130) | def compute_clip_score(
  function compute_fid (line 180) | def compute_fid(fake_dir: Path, gt_dir: Path,
  function evaluate_model (line 207) | def evaluate_model(opt):

FILE: language/extract_t5_feature.py
  class CustomDataset (line 23) | class CustomDataset(Dataset):
    method __init__ (line 24) | def __init__(self, lst_dir, start, end, caption_key, trunc_caption=Fal...
    method __len__ (line 41) | def __len__(self):
    method __getitem__ (line 44) | def __getitem__(self, index):
  function main (line 53) | def main(args):

FILE: language/t5.py
  class T5Embedder (line 15) | class T5Embedder:
    method __init__ (line 19) | def __init__(self, device, dir_or_name='t5-v1_1-xxl', *, local_cache=F...
    method get_text_embeddings (line 58) | def get_text_embeddings(self, texts):
    method text_preprocessing (line 81) | def text_preprocessing(self, text):
    method basic_clean (line 91) | def basic_clean(text):
    method clean_caption (line 96) | def clean_caption(self, caption):

FILE: tokenizer/consistencydecoder/cd_demo.py
  function main (line 9) | def main(args):

FILE: tokenizer/consistencydecoder/reconstruction_cd_ddp.py
  class SingleFolderDataset (line 22) | class SingleFolderDataset(Dataset):
    method __init__ (line 23) | def __init__(self, directory, transform=None):
    method __len__ (line 30) | def __len__(self):
    method __getitem__ (line 33) | def __getitem__(self, idx):
  function create_npz_from_sample_folder (line 41) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function center_crop_arr (line 60) | def center_crop_arr(pil_image, image_size):
  function main (line 81) | def main(args):

FILE: tokenizer/tokenizer_image/discriminator.py
  class PatchGANDiscriminator (line 17) | class PatchGANDiscriminator(nn.Module):
    method __init__ (line 21) | def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
    method _init_weights (line 67) | def _init_weights(self, module):
    method forward (line 74) | def forward(self, input):
  class ActNorm (line 79) | class ActNorm(nn.Module):
    method __init__ (line 80) | def __init__(self, num_features, logdet=False, affine=True,
    method initialize (line 91) | def initialize(self, input):
    method forward (line 112) | def forward(self, input, reverse=False):
    method reverse (line 140) | def reverse(self, output):
  class StyleGANDiscriminator (line 168) | class StyleGANDiscriminator(nn.Module):
    method __init__ (line 169) | def __init__(self, input_nc=3, ndf=64, n_layers=3, channel_multiplier=...
    method forward (line 203) | def forward(self, x):
  class DiscriminatorBlock (line 212) | class DiscriminatorBlock(nn.Module):
    method __init__ (line 213) | def __init__(self, input_channels, filters, downsample=True):
    method forward (line 229) | def forward(self, x):
  class Blur (line 238) | class Blur(nn.Module):
    method __init__ (line 239) | def __init__(self):
    method forward (line 244) | def forward(self, x):
  function leaky_relu (line 250) | def leaky_relu(p=0.2):
  function exists (line 254) | def exists(val):

FILE: tokenizer/tokenizer_image/discriminator_patchgan.py
  class NLayerDiscriminator (line 8) | class NLayerDiscriminator(nn.Module):
    method __init__ (line 12) | def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
    method _init_weights (line 58) | def _init_weights(self, module):
    method forward (line 65) | def forward(self, input):
  class ActNorm (line 70) | class ActNorm(nn.Module):
    method __init__ (line 71) | def __init__(self, num_features, logdet=False, affine=True,
    method initialize (line 82) | def initialize(self, input):
    method forward (line 103) | def forward(self, input, reverse=False):
    method reverse (line 131) | def reverse(self, output):

FILE: tokenizer/tokenizer_image/discriminator_stylegan.py
  class Discriminator (line 13) | class Discriminator(nn.Module):
    method __init__ (line 14) | def __init__(self, input_nc=3, ndf=64, n_layers=3, channel_multiplier=...
    method forward (line 48) | def forward(self, x):
  class DiscriminatorBlock (line 57) | class DiscriminatorBlock(nn.Module):
    method __init__ (line 58) | def __init__(self, input_channels, filters, downsample=True):
    method forward (line 74) | def forward(self, x):
  class Blur (line 84) | class Blur(nn.Module):
    method __init__ (line 85) | def __init__(self):
    method forward (line 90) | def forward(self, x):
  function leaky_relu (line 96) | def leaky_relu(p=0.2):
  function exists (line 100) | def exists(val):

FILE: tokenizer/tokenizer_image/lpips.py
  function download (line 24) | def download(url, local_path, chunk_size=1024):
  function md5_hash (line 36) | def md5_hash(path):
  function get_ckpt_path (line 42) | def get_ckpt_path(name, root, check=False):
  class LPIPS (line 53) | class LPIPS(nn.Module):
    method __init__ (line 55) | def __init__(self, use_dropout=True):
    method load_from_pretrained (line 69) | def load_from_pretrained(self, name="vgg_lpips"):
    method from_pretrained (line 75) | def from_pretrained(cls, name="vgg_lpips"):
    method forward (line 83) | def forward(self, input, target):
  class ScalingLayer (line 99) | class ScalingLayer(nn.Module):
    method __init__ (line 100) | def __init__(self):
    method forward (line 105) | def forward(self, inp):
  class NetLinLayer (line 109) | class NetLinLayer(nn.Module):
    method __init__ (line 111) | def __init__(self, chn_in, chn_out=1, use_dropout=False):
  class vgg16 (line 118) | class vgg16(torch.nn.Module):
    method __init__ (line 119) | def __init__(self, requires_grad=False, pretrained=True):
    method forward (line 142) | def forward(self, X):
  function normalize_tensor (line 158) | def normalize_tensor(x,eps=1e-10):
  function spatial_average (line 163) | def spatial_average(x, keepdim=True):

FILE: tokenizer/tokenizer_image/reconstruction_vq_ddp.py
  function create_npz_from_sample_folder (line 24) | def create_npz_from_sample_folder(sample_dir, num=50000):
  function main (line 42) | def main(args):

FILE: tokenizer/tokenizer_image/vq_demo.py
  function main (line 13) | def main(args):

FILE: tokenizer/tokenizer_image/vq_loss.py
  function hinge_d_loss (line 14) | def hinge_d_loss(logits_real, logits_fake):
  function vanilla_d_loss (line 21) | def vanilla_d_loss(logits_real, logits_fake):
  function non_saturating_d_loss (line 28) | def non_saturating_d_loss(logits_real, logits_fake):
  function hinge_gen_loss (line 35) | def hinge_gen_loss(logit_fake):
  function non_saturating_gen_loss (line 39) | def non_saturating_gen_loss(logit_fake):
  function adopt_weight (line 43) | def adopt_weight(weight, global_step, threshold=0, value=0.):
  class VQLoss (line 49) | class VQLoss(nn.Module):
    method __init__ (line 50) | def __init__(self, disc_start, disc_loss="hinge", disc_dim=64, disc_ty...
    method calculate_adaptive_weight (line 109) | def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer):
    method forward (line 117) | def forward(self, codebook_loss, inputs, reconstructions, optimizer_id...

FILE: tokenizer/tokenizer_image/vq_model.py
  class ModelArgs (line 13) | class ModelArgs:
  class VQModel (line 28) | class VQModel(nn.Module):
    method __init__ (line 29) | def __init__(self, config: ModelArgs):
    method encode (line 41) | def encode(self, x):
    method decode (line 48) | def decode(self, quant):
    method decode_code (line 53) | def decode_code(self, code_b, shape=None, channel_first=True):
    method forward (line 58) | def forward(self, input):
  class Encoder (line 65) | class Encoder(nn.Module):
    method __init__ (line 66) | def __init__(self, in_channels=3, ch=128, ch_mult=(1,1,2,2,4), num_res...
    method forward (line 106) | def forward(self, x):
  class Decoder (line 129) | class Decoder(nn.Module):
    method __init__ (line 130) | def __init__(self, z_channels=256, ch=128, ch_mult=(1,1,2,2,4), num_re...
    method last_layer (line 171) | def last_layer(self):
    method forward (line 174) | def forward(self, z):
  class VectorQuantizer (line 198) | class VectorQuantizer(nn.Module):
    method __init__ (line 199) | def __init__(self, n_e, e_dim, beta, entropy_loss_ratio, l2_norm, show...
    method forward (line 216) | def forward(self, z):
    method get_codebook_entry (line 262) | def get_codebook_entry(self, indices, shape=None, channel_first=True):
  class ResnetBlock (line 280) | class ResnetBlock(nn.Module):
    method __init__ (line 281) | def __init__(self, in_channels, out_channels=None, conv_shortcut=False...
    method forward (line 300) | def forward(self, x):
  class AttnBlock (line 318) | class AttnBlock(nn.Module):
    method __init__ (line 319) | def __init__(self, in_channels, norm_type='group'):
    method forward (line 328) | def forward(self, x):
  function nonlinearity (line 355) | def nonlinearity(x):
  function Normalize (line 360) | def Normalize(in_channels, norm_type='group'):
  class Upsample (line 368) | class Upsample(nn.Module):
    method __init__ (line 369) | def __init__(self, in_channels, with_conv):
    method forward (line 375) | def forward(self, x):
  class Downsample (line 382) | class Downsample(nn.Module):
    method __init__ (line 383) | def __init__(self, in_channels, with_conv):
    method forward (line 390) | def forward(self, x):
  function compute_entropy_loss (line 400) | def compute_entropy_loss(affinity, loss_type="softmax", temperature=0.01):
  function VQ_8 (line 419) | def VQ_8(**kwargs):
  function VQ_16 (line 422) | def VQ_16(**kwargs):

FILE: tokenizer/tokenizer_image/vq_model_hf.py
  class VQModelHF (line 5) | class VQModelHF(VQModel, PyTorchModelHubMixin, repo_url="https://github....
  function VQ_8 (line 11) | def VQ_8(**kwargs):
  function VQ_16 (line 14) | def VQ_16(**kwargs):

FILE: tokenizer/tokenizer_image/vq_train.py
  function main (line 37) | def main(args):

FILE: tokenizer/vae/reconstruction_vae_ddp.py
  class SingleFolderDataset (line 22) | class SingleFolderDataset(Dataset):
    method __init__ (line 23) | def __init__(self, directory, transform=None):
    method __len__ (line 30) | def __len__(self):
    method __getitem__ (line 33) | def __getitem__(self, idx):
  function create_npz_from_sample_folder (line 41) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function center_crop_arr (line 60) | def center_crop_arr(pil_image, image_size):
  function main (line 81) | def main(args):

FILE: tokenizer/vae/sd_vae_demo.py
  function main (line 9) | def main(args):

FILE: tokenizer/validation/val_ddp.py
  class SingleFolderDataset (line 17) | class SingleFolderDataset(Dataset):
    method __init__ (line 18) | def __init__(self, directory, transform=None):
    method __len__ (line 25) | def __len__(self):
    method __getitem__ (line 28) | def __getitem__(self, idx):
  function create_npz_from_sample_folder (line 36) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function center_crop_arr (line 55) | def center_crop_arr(pil_image, image_size):
  function main (line 76) | def main(args):

FILE: tokenizer/vqgan/layer.py
  function nonlinearity (line 8) | def nonlinearity(x):
  function Normalize (line 13) | def Normalize(in_channels):
  class Upsample (line 17) | class Upsample(nn.Module):
    method __init__ (line 18) | def __init__(self, in_channels, with_conv):
    method forward (line 28) | def forward(self, x):
  class Downsample (line 35) | class Downsample(nn.Module):
    method __init__ (line 36) | def __init__(self, in_channels, with_conv):
    method forward (line 47) | def forward(self, x):
  class ResnetBlock (line 57) | class ResnetBlock(nn.Module):
    method __init__ (line 58) | def __init__(self, *, in_channels, out_channels=None, conv_shortcut=Fa...
    method forward (line 96) | def forward(self, x, temb):
  class AttnBlock (line 119) | class AttnBlock(nn.Module):
    method __init__ (line 120) | def __init__(self, in_channels):
    method forward (line 147) | def forward(self, x):
  class Encoder (line 175) | class Encoder(nn.Module):
    method __init__ (line 176) | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
    method forward (line 239) | def forward(self, x):
  class Decoder (line 269) | class Decoder(nn.Module):
    method __init__ (line 270) | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
    method forward (line 339) | def forward(self, z):

FILE: tokenizer/vqgan/model.py
  class VQModel (line 24) | class VQModel(nn.Module):
    method __init__ (line 25) | def __init__(self,
    method init_from_ckpt (line 55) | def init_from_ckpt(self, path, ignore_keys=list(), logging=True):
    method encode (line 69) | def encode(self, x):
    method decode (line 75) | def decode(self, quant):
    method decode_code (line 80) | def decode_code(self, code_b, shape, channel_first=True):
    method forward (line 85) | def forward(self, input):

FILE: tokenizer/vqgan/quantize.py
  class VectorQuantizer (line 9) | class VectorQuantizer(nn.Module):
    method __init__ (line 25) | def __init__(self, n_e, e_dim, beta):
    method forward (line 34) | def forward(self, z):
    method get_codebook_entry (line 92) | def get_codebook_entry(self, indices, shape):
  class VectorQuantizer2 (line 110) | class VectorQuantizer2(nn.Module):
    method __init__ (line 118) | def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random",
    method remap_to_used (line 144) | def remap_to_used(self, inds):
    method unmap_to_all (line 158) | def unmap_to_all(self, inds):
    method forward (line 168) | def forward(self, z, temp=None, rescale_logits=False, return_logits=Fa...
    method get_codebook_entry (line 211) | def get_codebook_entry(self, indices, shape, channel_first=True):

FILE: tokenizer/vqgan/reconstruction_vqgan_ddp.py
  class SingleFolderDataset (line 24) | class SingleFolderDataset(Dataset):
    method __init__ (line 25) | def __init__(self, directory, transform=None):
    method __len__ (line 32) | def __len__(self):
    method __getitem__ (line 35) | def __getitem__(self, idx):
  function create_npz_from_sample_folder (line 43) | def create_npz_from_sample_folder(sample_dir, num=50_000):
  function center_crop_arr (line 62) | def center_crop_arr(pil_image, image_size):
  function main (line 83) | def main(args):

FILE: tokenizer/vqgan/taming_vqgan_demo.py
  function main (line 17) | def main(args):

FILE: tools/check_image_codes.py
  function main (line 9) | def main(args):

FILE: tools/draw_figure.py
  function fid_scaling_law_no_cfg (line 6) | def fid_scaling_law_no_cfg():
  function fid_scaling_law_cfg (line 43) | def fid_scaling_law_cfg():
  function sample_topk (line 80) | def sample_topk():
  function sample_cfg (line 108) | def sample_cfg():

FILE: tools/openimage_json.py
  function check_image (line 11) | def check_image(image_path):
  function check_image_path (line 20) | def check_image_path(image_info):
  function load_image_path (line 29) | def load_image_path(image_info):
  function main (line 44) | def main(args):

FILE: tools/push_gpt_to_hf.py
  function main (line 13) | def main(args):

FILE: tools/push_vae_to_hf.py
  function load_model (line 17) | def load_model(args):

FILE: utils/data.py
  function center_crop_arr (line 4) | def center_crop_arr(pil_image, image_size):

FILE: utils/deepspeed.py
  function create_deepspeed_config (line 1) | def create_deepspeed_config(args):

FILE: utils/distributed.py
  function setup_for_distributed (line 6) | def setup_for_distributed(is_master):
  function init_distributed_mode (line 20) | def init_distributed_mode(args):

FILE: utils/drop_path.py
  function drop_path (line 4) | def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by...
  class DropPath (line 24) | class DropPath(torch.nn.Module):
    method __init__ (line 27) | def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
    method forward (line 32) | def forward(self, x):
    method extra_repr (line 35) | def extra_repr(self):

FILE: utils/ema.py
  function update_ema (line 5) | def update_ema(ema_model, model, decay=0.9999):
  function requires_grad (line 17) | def requires_grad(model, flag=True):

FILE: utils/logger.py
  function create_logger (line 4) | def create_logger(logging_dir):

FILE: utils/video.py
  function shift_dim (line 8) | def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
  function view_range (line 38) | def view_range(x, i, j, shape):
  function tensor_slice (line 57) | def tensor_slice(x, begin, size):
  function save_video_grid (line 67) | def save_video_grid(video, fname, nrow=None, fps=5):
  function save_gif_grid (line 89) | def save_gif_grid(video, file_name, nrow=None, fps=5):