SYMBOL INDEX (598 symbols across 55 files)

FILE: convert_checkpoint.py
  function main (line 7) | def main():

FILE: demo.py
  function initialize_vae_decoder (line 53) | def initialize_vae_decoder(use_taehv=False, use_trt=False):
  function tensor_to_base64_frame (line 155) | def tensor_to_base64_frame(frame_tensor):
  function frame_sender_worker (line 178) | def frame_sender_worker():
  function generate_video_stream (line 230) | def generate_video_stream(prompt, seed, enable_torch_compile=False, enab...
  function handle_connect (line 499) | def handle_connect():
  function handle_disconnect (line 505) | def handle_disconnect():
  function handle_start_generation (line 510) | def handle_start_generation(data):
  function handle_stop_generation (line 534) | def handle_stop_generation():
  function index (line 551) | def index():
  function api_status (line 556) | def api_status():

FILE: demo_utils/memory.py
  class DynamicSwapInstaller (line 13) | class DynamicSwapInstaller:
    method _install_module (line 15) | def _install_module(module: torch.nn.Module, **kwargs):
    method _uninstall_module (line 43) | def _uninstall_module(module: torch.nn.Module):
    method install_model (line 49) | def install_model(model: torch.nn.Module, **kwargs):
    method uninstall_model (line 55) | def uninstall_model(model: torch.nn.Module):
  function fake_diffusers_current_device (line 61) | def fake_diffusers_current_device(model: torch.nn.Module, target_device:...
  function get_cuda_free_memory_gb (line 72) | def get_cuda_free_memory_gb(device=None):
  function move_model_to_device_with_memory_preservation (line 85) | def move_model_to_device_with_memory_preservation(model, target_device, ...
  function offload_model_from_device_for_memory_preservation (line 101) | def offload_model_from_device_for_memory_preservation(model, target_devi...
  function unload_complete_models (line 117) | def unload_complete_models(*args):
  function load_model_as_complete (line 127) | def load_model_as_complete(model, target_device, unload=True):

FILE: demo_utils/taehv.py
  function conv (line 16) | def conv(n_in, n_out, **kwargs):
  class Clamp (line 20) | class Clamp(nn.Module):
    method forward (line 21) | def forward(self, x):
  class MemBlock (line 25) | class MemBlock(nn.Module):
    method __init__ (line 26) | def __init__(self, n_in, n_out):
    method forward (line 33) | def forward(self, x, past):
  class TPool (line 37) | class TPool(nn.Module):
    method __init__ (line 38) | def __init__(self, n_f, stride):
    method forward (line 43) | def forward(self, x):
  class TGrow (line 48) | class TGrow(nn.Module):
    method __init__ (line 49) | def __init__(self, n_f, stride):
    method forward (line 54) | def forward(self, x):
  function apply_model_with_memblocks (line 60) | def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
  class TAEHV (line 159) | class TAEHV(nn.Module):
    method __init__ (line 163) | def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(...
    method patch_tgrow_layers (line 195) | def patch_tgrow_layers(self, sd):
    method encode_video (line 210) | def encode_video(self, x, parallel=True, show_progress_bar=True):
    method decode_video (line 222) | def decode_video(self, x, parallel=True, show_progress_bar=False):
    method forward (line 236) | def forward(self, x):
  function main (line 241) | def main():

FILE: demo_utils/utils.py
  function min_resize (line 19) | def min_resize(x, m):
  function d_resize (line 36) | def d_resize(x, y):
  function resize_and_center_crop (line 48) | def resize_and_center_crop(image, target_width, target_height):
  function resize_and_center_crop_pytorch (line 66) | def resize_and_center_crop_pytorch(image, target_width, target_height):
  function resize_without_crop (line 85) | def resize_without_crop(image, target_width, target_height):
  function just_crop (line 94) | def just_crop(image, w, h):
  function write_to_json (line 108) | def write_to_json(data, file_path):
  function read_from_json (line 116) | def read_from_json(file_path):
  function get_active_parameters (line 122) | def get_active_parameters(m):
  function cast_training_params (line 126) | def cast_training_params(m, dtype=torch.float32):
  function separate_lora_AB (line 135) | def separate_lora_AB(parameters, B_patterns=None):
  function set_attr_recursive (line 151) | def set_attr_recursive(obj, attr, value):
  function print_tensor_list_size (line 159) | def print_tensor_list_size(tensors):
  function batch_mixture (line 180) | def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
  function zero_module (line 196) | def zero_module(module):
  function supress_lower_channels (line 203) | def supress_lower_channels(m, k, alpha=0.01):
  function freeze_module (line 213) | def freeze_module(m):
  function get_latest_safetensors (line 221) | def get_latest_safetensors(folder_path):
  function generate_random_prompt_from_tags (line 232) | def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=...
  function interpolate_numbers (line 239) | def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
  function uniform_random_by_intervals (line 246) | def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=Fa...
  function soft_append_bcthw (line 255) | def soft_append_bcthw(history, current, overlap=0):
  function save_bcthw_as_mp4 (line 269) | def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
  function save_bcthw_as_png (line 286) | def save_bcthw_as_png(x, output_filename):
  function save_bchw_as_png (line 295) | def save_bchw_as_png(x, output_filename):
  function add_tensors_with_padding (line 304) | def add_tensors_with_padding(tensor1, tensor2):
  function print_free_mem (line 323) | def print_free_mem():
  function print_gpu_parameters (line 333) | def print_gpu_parameters(device, state_dict, log_count=1):
  function visualize_txt_as_img (line 348) | def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans...
  function blue_mark (line 386) | def blue_mark(x):
  function green_mark (line 394) | def green_mark(x):
  function frame_mark (line 401) | def frame_mark(x):
  function pytorch2numpy (line 411) | def pytorch2numpy(imgs):
  function numpy2pytorch (line 422) | def numpy2pytorch(imgs):
  function duplicate_prefix_to_suffix (line 429) | def duplicate_prefix_to_suffix(x, count, zero_out=False):
  function weighted_mse (line 436) | def weighted_mse(a, b, weight):
  function clamped_linear_interpolation (line 440) | def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
  function expand_to_dims (line 447) | def expand_to_dims(x, target_dims):
  function repeat_to_batch_size (line 451) | def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
  function dim5 (line 468) | def dim5(x):
  function dim4 (line 472) | def dim4(x):
  function dim3 (line 476) | def dim3(x):
  function crop_or_pad_yield_mask (line 480) | def crop_or_pad_yield_mask(x, length):
  function extend_dim (line 495) | def extend_dim(x, dim, minimal_length, zero_pad=False):
  function lazy_positional_encoding (line 513) | def lazy_positional_encoding(t, repeats=None):
  function state_dict_offset_merge (line 530) | def state_dict_offset_merge(A, B, C=None):
  function state_dict_weighted_merge (line 547) | def state_dict_weighted_merge(state_dicts, weights):
  function group_files_by_folder (line 574) | def group_files_by_folder(all_files):
  function generate_timestamp (line 587) | def generate_timestamp():
  function write_PIL_image_with_png_info (line 595) | def write_PIL_image_with_png_info(image, metadata, path):
  function torch_safe_save (line 606) | def torch_safe_save(content, path):
  function move_optimizer_to_device (line 612) | def move_optimizer_to_device(optimizer, device):

FILE: demo_utils/vae.py
  class ResidualBlock (line 13) | class ResidualBlock(nn.Module):
    method __init__ (line 15) | def __init__(self, in_dim, out_dim, dropout=0.0):
    method forward (line 29) | def forward(self, x, feat_cache_1, feat_cache_2):
  class Resample (line 51) | class Resample(nn.Module):
    method __init__ (line 53) | def __init__(self, dim, mode):
    method forward (line 73) | def forward(self, x, is_first_frame, feat_cache):
    method temporal_conv (line 105) | def temporal_conv(self, x, is_first_frame, feat_cache):
    method init_weight (line 127) | def init_weight(self, conv):
    method init_weight2 (line 139) | def init_weight2(self, conv):
  class VAEDecoderWrapperSingle (line 151) | class VAEDecoderWrapperSingle(nn.Module):
    method __init__ (line 152) | def __init__(self):
    method forward (line 168) | def forward(
  class VAEDecoder3d (line 199) | class VAEDecoder3d(nn.Module):
    method __init__ (line 200) | def __init__(self,
    method forward (line 254) | def forward(
  class VAETRTWrapper (line 318) | class VAETRTWrapper():
    method __init__ (line 319) | def __init__(self):
    method quantize_if_needed (line 376) | def quantize_if_needed(self, t, expected_dtype, scale):
    method forward (line 381) | def forward(self, *test_inputs):

FILE: demo_utils/vae_block3.py
  class Resample (line 9) | class Resample(nn.Module):
    method __init__ (line 11) | def __init__(self, dim, mode):
    method forward (line 45) | def forward(self, x, feat_cache=None, feat_idx=[0]):
    method init_weight (line 106) | def init_weight(self, conv):
    method init_weight2 (line 118) | def init_weight2(self, conv):
  class VAEDecoderWrapper (line 130) | class VAEDecoderWrapper(nn.Module):
    method __init__ (line 131) | def __init__(self):
    method forward (line 147) | def forward(
  class VAEDecoder3d (line 187) | class VAEDecoder3d(nn.Module):
    method __init__ (line 188) | def __init__(self,
    method forward (line 242) | def forward(

FILE: demo_utils/vae_torch2trt.py
  function set_workspace (line 98) | def set_workspace(config, bytes_):
  function set_workspace (line 122) | def set_workspace(config: trt.IBuilderConfig, bytes_: int = 4 << 30):
  class VAECalibrator (line 139) | class VAECalibrator(trt.IInt8EntropyCalibrator2):
    method __init__ (line 140) | def __init__(self, loader, cache="calibration.cache", max_batches=10):
    method get_batch_size (line 151) | def get_batch_size(self):
    method getBatchSize (line 154) | def getBatchSize(self):
    method get_batch (line 157) | def get_batch(self, names):
    method read_calibration_cache (line 202) | def read_calibration_cache(self):
    method readCalibrationCache (line 209) | def readCalibrationCache(self):
    method write_calibration_cache (line 212) | def write_calibration_cache(self, cache):
    method writeCalibrationCache (line 216) | def writeCalibrationCache(self, cache):

FILE: inference.py
  function encode (line 98) | def encode(self, videos: torch.Tensor) -> torch.Tensor:

FILE: model/base.py
  class BaseModel (line 12) | class BaseModel(nn.Module):
    method __init__ (line 13) | def __init__(self, args, device):
    method _initialize_models (line 28) | def _initialize_models(self, args, device):
    method _get_timestep (line 59) | def _get_timestep(
  class SelfForcingModel (line 109) | class SelfForcingModel(BaseModel):
    method __init__ (line 110) | def __init__(self, args, device):
    method _run_generator (line 114) | def _run_generator(
    method _consistency_backward_simulation (line 197) | def _consistency_backward_simulation(
    method _initialize_inference_pipeline (line 223) | def _initialize_inference_pipeline(self):

FILE: model/causvid.py
  class CausVid (line 8) | class CausVid(BaseModel):
    method __init__ (line 9) | def __init__(self, args, device):
    method _compute_kl_grad (line 47) | def _compute_kl_grad(
    method compute_distribution_matching_loss (line 121) | def compute_distribution_matching_loss(
    method _run_generator (line 184) | def _run_generator(
    method generator_loss (line 255) | def generator_loss(
    method critic_loss (line 296) | def critic_loss(

FILE: model/diffusion.py
  class CausalDiffusion (line 8) | class CausalDiffusion(BaseModel):
    method __init__ (line 9) | def __init__(self, args, device):
    method _initialize_models (line 34) | def _initialize_models(self, args):
    method generator_loss (line 44) | def generator_loss(

FILE: model/dmd.py
  class DMD (line 9) | class DMD(SelfForcingModel):
    method __init__ (line 10) | def __init__(self, args, device):
    method _compute_kl_grad (line 54) | def _compute_kl_grad(
    method compute_distribution_matching_loss (line 138) | def compute_distribution_matching_loss(
    method generator_loss (line 210) | def generator_loss(
    method critic_loss (line 259) | def critic_loss(

FILE: model/gan.py
  class GAN (line 10) | class GAN(SelfForcingModel):
    method __init__ (line 11) | def __init__(self, args, device):
    method _run_cls_pred_branch (line 69) | def _run_cls_pred_branch(self,
    method generator_loss (line 90) | def generator_loss(
    method critic_loss (line 174) | def critic_loss(

FILE: model/ode_regression.py
  class ODERegression (line 9) | class ODERegression(BaseModel):
    method __init__ (line 10) | def __init__(self, args, device):
    method _initialize_models (line 46) | def _initialize_models(self, args):
    method _prepare_generator_input (line 57) | def _prepare_generator_input(self, ode_latent: torch.Tensor) -> Tuple[...
    method generator_loss (line 102) | def generator_loss(self, ode_latent: torch.Tensor, conditional_dict: d...

FILE: model/sid.py
  class SiD (line 8) | class SiD(SelfForcingModel):
    method __init__ (line 9) | def __init__(self, args, device):
    method compute_distribution_matching_loss (line 47) | def compute_distribution_matching_loss(
    method generator_loss (line 147) | def generator_loss(
    method critic_loss (line 188) | def critic_loss(

FILE: pipeline/bidirectional_diffusion_inference.py
  class BidirectionalDiffusionInferencePipeline (line 10) | class BidirectionalDiffusionInferencePipeline(torch.nn.Module):
    method __init__ (line 11) | def __init__(
    method inference (line 34) | def inference(
    method _initialize_sample_scheduler (line 89) | def _initialize_sample_scheduler(self, noise):

FILE: pipeline/bidirectional_inference.py
  class BidirectionalInferencePipeline (line 7) | class BidirectionalInferencePipeline(torch.nn.Module):
    method __init__ (line 8) | def __init__(
    method inference (line 33) | def inference(self, noise: torch.Tensor, text_prompts: List[str]) -> t...

FILE: pipeline/bidirectional_training.py
  class BidirectionalTrainingPipeline (line 9) | class BidirectionalTrainingPipeline(torch.nn.Module):
    method __init__ (line 10) | def __init__(
    method generate_and_sync_list (line 25) | def generate_and_sync_list(self, num_denoising_steps, device):
    method inference_with_trajectory (line 42) | def inference_with_trajectory(self, noise: torch.Tensor, clip_fea, y, ...

FILE: pipeline/causal_diffusion_inference.py
  class CausalDiffusionInferencePipeline (line 10) | class CausalDiffusionInferencePipeline(torch.nn.Module):
    method __init__ (line 11) | def __init__(
    method inference (line 49) | def inference(
    method _initialize_kv_cache (line 270) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 300) | def _initialize_crossattn_cache(self, batch_size, dtype, device):
    method _initialize_sample_scheduler (line 321) | def _initialize_sample_scheduler(self, noise):

FILE: pipeline/causal_inference.py
  class CausalInferencePipeline (line 7) | class CausalInferencePipeline(torch.nn.Module):
    method __init__ (line 8) | def __init__(
    method inference (line 45) | def inference(
    method _initialize_kv_cache (line 271) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 293) | def _initialize_crossattn_cache(self, batch_size, dtype, device):

FILE: pipeline/self_forcing_training.py
  class SelfForcingTrainingPipeline (line 8) | class SelfForcingTrainingPipeline:
    method __init__ (line 9) | def __init__(self,
    method generate_and_sync_list (line 45) | def generate_and_sync_list(self, num_blocks, num_denoising_steps, devi...
    method inference_with_trajectory (line 64) | def inference_with_trajectory(
    method _initialize_kv_cache (line 245) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 261) | def _initialize_crossattn_cache(self, batch_size, dtype, device):

FILE: scripts/compute_vae_latent.py
  function launch_distributed_job (line 16) | def launch_distributed_job(backend: str = "nccl"):
  function video_to_numpy (line 32) | def video_to_numpy(video_path):
  function encode (line 42) | def encode(self, videos: torch.Tensor) -> torch.Tensor:
  function main (line 55) | def main():

FILE: scripts/create_lmdb_14b_shards.py
  function main (line 19) | def main():

FILE: scripts/create_lmdb_iterative.py
  function main (line 12) | def main():

FILE: scripts/generate_ode_pairs.py
  function init_model (line 13) | def init_model(device):
  function main (line 32) | def main():

FILE: train.py
  function main (line 9) | def main():

FILE: trainer/diffusion.py
  class Trainer (line 17) | class Trainer:
    method __init__ (line 18) | def __init__(self, config):
    method save (line 140) | def save(self):
    method train_one_step (line 163) | def train_one_step(self, batch):
    method generate_video (line 235) | def generate_video(self, pipeline, prompts, image=None):
    method train (line 248) | def train(self):

FILE: trainer/distillation.py
  class Trainer (line 20) | class Trainer:
    method __init__ (line 21) | def __init__(self, config):
    method save (line 236) | def save(self):
    method fwdbwd_one_step (line 263) | def fwdbwd_one_step(self, batch, train_generator):
    method generate_video (line 350) | def generate_video(self, pipeline, prompts, image=None):
    method train (line 380) | def train(self):

FILE: trainer/gan.py
  class Trainer (line 19) | class Trainer:
    method __init__ (line 20) | def __init__(self, config):
    method save (line 208) | def save(self):
    method fwdbwd_one_step (line 235) | def fwdbwd_one_step(self, batch, train_generator):
    method generate_video (line 324) | def generate_video(self, pipeline, prompts, image=None):
    method train (line 337) | def train(self):
    method all_gather_dict (line 457) | def all_gather_dict(self, target_dict):

FILE: trainer/ode.py
  class Trainer (line 19) | class Trainer:
    method __init__ (line 20) | def __init__(self, config):
    method save (line 118) | def save(self):
    method train_one_step (line 134) | def train_one_step(self):
    method train (line 225) | def train(self):

FILE: utils/dataset.py
  class TextDataset (line 13) | class TextDataset(Dataset):
    method __init__ (line 14) | def __init__(self, prompt_path, extended_prompt_path=None):
    method __len__ (line 25) | def __len__(self):
    method __getitem__ (line 28) | def __getitem__(self, idx):
  class TextFolderDataset (line 38) | class TextFolderDataset(Dataset):
    method __init__ (line 39) | def __init__(self, data_path, max_count=30000):
    method __len__ (line 51) | def __len__(self):
    method __getitem__ (line 54) | def __getitem__(self, idx):
  class ODERegressionLMDBDataset (line 58) | class ODERegressionLMDBDataset(Dataset):
    method __init__ (line 59) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 66) | def __len__(self):
    method __getitem__ (line 69) | def __getitem__(self, idx):
  class ShardingLMDBDataset (line 93) | class ShardingLMDBDataset(Dataset):
    method __init__ (line 94) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 117) | def __len__(self):
    method __getitem__ (line 120) | def __getitem__(self, idx):
  class TextImagePairDataset (line 157) | class TextImagePairDataset(Dataset):
    method __init__ (line 158) | def __init__(
    method __len__ (line 212) | def __len__(self):
    method __getitem__ (line 215) | def __getitem__(self, idx):
  function cycle (line 247) | def cycle(dl):

FILE: utils/distributed.py
  function fsdp_state_dict (line 11) | def fsdp_state_dict(model):
  function fsdp_wrap (line 23) | def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, w...
  function barrier (line 71) | def barrier():
  function launch_distributed_job (line 76) | def launch_distributed_job(backend: str = "nccl"):
  class EMA_FSDP (line 92) | class EMA_FSDP:
    method __init__ (line 93) | def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
    method _init_shadow (line 99) | def _init_shadow(self, fsdp_module):
    method update (line 106) | def update(self, fsdp_module):
    method state_dict (line 114) | def state_dict(self):
    method load_state_dict (line 117) | def load_state_dict(self, sd):
    method copy_to (line 120) | def copy_to(self, fsdp_module):

FILE: utils/lmdb.py
  function get_array_shape_from_lmdb (line 4) | def get_array_shape_from_lmdb(env, array_name):
  function store_arrays_to_lmdb (line 11) | def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
  function process_data_dict (line 30) | def process_data_dict(data_dict, seen_prompts):
  function retrieve_row_from_lmdb (line 57) | def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape...

FILE: utils/loss.py
  class DenoisingLoss (line 5) | class DenoisingLoss(ABC):
    method __call__ (line 7) | def __call__(
  class X0PredLoss (line 27) | class X0PredLoss(DenoisingLoss):
    method __call__ (line 28) | def __call__(
  class VPredLoss (line 38) | class VPredLoss(DenoisingLoss):
    method __call__ (line 39) | def __call__(
  class NoisePredLoss (line 50) | class NoisePredLoss(DenoisingLoss):
    method __call__ (line 51) | def __call__(
  class FlowPredLoss (line 61) | class FlowPredLoss(DenoisingLoss):
    method __call__ (line 62) | def __call__(
  function get_denoising_loss (line 80) | def get_denoising_loss(loss_type: str) -> DenoisingLoss:

FILE: utils/misc.py
  function set_seed (line 6) | def set_seed(seed: int, deterministic: bool = False):
  function merge_dict_list (line 25) | def merge_dict_list(dict_list):

FILE: utils/scheduler.py
  class SchedulerInterface (line 5) | class SchedulerInterface(ABC):
    method add_noise (line 12) | def add_noise(
    method convert_x0_to_noise (line 26) | def convert_x0_to_noise(
    method convert_noise_to_x0 (line 52) | def convert_noise_to_x0(
    method convert_velocity_to_x0 (line 77) | def convert_velocity_to_x0(
  class FlowMatchScheduler (line 106) | class FlowMatchScheduler():
    method __init__ (line 108) | def __init__(self, num_inference_steps=100, num_train_timesteps=1000, ...
    method set_timesteps (line 118) | def set_timesteps(self, num_inference_steps=100, denoising_strength=1....
    method step (line 143) | def step(self, model_output, timestep, sample, to_final=False):
    method add_noise (line 159) | def add_noise(self, original_samples, noise, timestep):
    method training_target (line 178) | def training_target(self, sample, noise, timestep):
    method training_weight (line 182) | def training_weight(self, timestep):

FILE: utils/wan_wrapper.py
  class WanTextEncoder (line 16) | class WanTextEncoder(torch.nn.Module):
    method __init__ (line 17) | def __init__(self, model_name="Wan2.1-T2V-14B") -> None:
    method device (line 36) | def device(self):
    method forward (line 40) | def forward(self, text_prompts: List[str]) -> dict:
  class WanCLIPEncoder (line 56) | class WanCLIPEncoder(torch.nn.Module):
    method __init__ (line 57) | def __init__(self, model_name="Wan2.1-T2V-14B"):
    method device (line 70) | def device(self):
    method forward (line 74) | def forward(self, img):
  class WanVAEWrapper (line 81) | class WanVAEWrapper(torch.nn.Module):
    method __init__ (line 82) | def __init__(self, model_name="Wan2.1-T2V-14B"):
    method encode (line 107) | def encode(self, pixel):
    method run_vae_encoder (line 117) | def run_vae_encoder(self, img):
    method encode_to_latent (line 149) | def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
    method decode_to_pixel (line 165) | def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = Fals...
  class WanDiffusionWrapper (line 191) | class WanDiffusionWrapper(torch.nn.Module):
    method __init__ (line 192) | def __init__(
    method enable_gradient_checkpointing (line 222) | def enable_gradient_checkpointing(self) -> None:
    method adding_cls_branch (line 225) | def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_di...
    method _convert_flow_pred_to_x0 (line 248) | def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch....
    method _convert_x0_to_flow_pred (line 275) | def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: tor...
    method forward (line 297) | def forward(
    method get_scheduler (line 380) | def get_scheduler(self) -> SchedulerInterface:
    method post_init (line 394) | def post_init(self):

FILE: wan/distributed/fsdp.py
  function shard_model (line 10) | def shard_model(

FILE: wan/distributed/xdit_context_parallel.py
  function pad_freqs (line 12) | def pad_freqs(original_tensor, target_len):
  function rope_apply (line 26) | def rope_apply(x, grid_sizes, freqs):
  function usp_dit_forward (line 66) | def usp_dit_forward(
  function usp_attn_forward (line 149) | def usp_attn_forward(self,

FILE: wan/image2video.py
  class WanI2V (line 29) | class WanI2V:
    method __init__ (line 31) | def __init__(
    method generate (line 129) | def generate(self,

FILE: wan/modules/attention.py
  function is_hopper_gpu (line 7) | def is_hopper_gpu():
  function flash_attention (line 32) | def flash_attention(
  function attention (line 139) | def attention(

FILE: wan/modules/causal_model.py
  function causal_rope_apply (line 27) | def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
  class CausalWanSelfAttention (line 58) | class CausalWanSelfAttention(nn.Module):
    method __init__ (line 60) | def __init__(self,
    method forward (line 86) | def forward(
  class CausalWanAttentionBlock (line 243) | class CausalWanAttentionBlock(nn.Module):
    method __init__ (line 245) | def __init__(self,
    method forward (line 283) | def forward(
  class CausalHead (line 338) | class CausalHead(nn.Module):
    method __init__ (line 340) | def __init__(self, dim, out_dim, patch_size, eps=1e-6):
    method forward (line 355) | def forward(self, x, e):
  class CausalWanModel (line 369) | class CausalWanModel(ModelMixin, ConfigMixin):
    method __init__ (line 381) | def __init__(self,
    method _set_gradient_checkpointing (line 502) | def _set_gradient_checkpointing(self, module, value=False):
    method _prepare_blockwise_causal_attn_mask (line 506) | def _prepare_blockwise_causal_attn_mask(
    method _prepare_teacher_forcing_mask (line 564) | def _prepare_teacher_forcing_mask(
    method _prepare_blockwise_causal_attn_mask_i2v (line 652) | def _prepare_blockwise_causal_attn_mask_i2v(
    method _forward_inference (line 712) | def _forward_inference(
    method _forward_train (line 843) | def _forward_train(
    method forward (line 1001) | def forward(
    method unpatchify (line 1011) | def unpatchify(self, x, grid_sizes):
    method init_weights (line 1036) | def init_weights(self):

FILE: wan/modules/clip.py
  function pos_interpolate (line 22) | def pos_interpolate(pos, seq_len):
  class QuickGELU (line 41) | class QuickGELU(nn.Module):
    method forward (line 43) | def forward(self, x):
  class LayerNorm (line 47) | class LayerNorm(nn.LayerNorm):
    method forward (line 49) | def forward(self, x):
  class SelfAttention (line 54) | class SelfAttention(nn.Module):
    method __init__ (line 56) | def __init__(self,
    method forward (line 75) | def forward(self, x):
  class SwiGLU (line 95) | class SwiGLU(nn.Module):
    method __init__ (line 97) | def __init__(self, dim, mid_dim):
    method forward (line 107) | def forward(self, x):
  class AttentionBlock (line 113) | class AttentionBlock(nn.Module):
    method __init__ (line 115) | def __init__(self,
    method forward (line 147) | def forward(self, x):
  class AttentionPool (line 157) | class AttentionPool(nn.Module):
    method __init__ (line 159) | def __init__(self,
    method forward (line 187) | def forward(self, x):
  class VisionTransformer (line 210) | class VisionTransformer(nn.Module):
    method __init__ (line 212) | def __init__(self,
    method forward (line 280) | def forward(self, x, interpolation=False, use_31_block=False):
  class XLMRobertaWithHead (line 304) | class XLMRobertaWithHead(XLMRoberta):
    method __init__ (line 306) | def __init__(self, **kwargs):
    method forward (line 316) | def forward(self, ids):
  class XLMRobertaCLIP (line 329) | class XLMRobertaCLIP(nn.Module):
    method __init__ (line 331) | def __init__(self,
  function _clip (line 408) | def _clip(pretrained=False,
  function clip_xlm_roberta_vit_h_14 (line 445) | def clip_xlm_roberta_vit_h_14(
  class CLIPModel (line 475) | class CLIPModel(nn.Module):
    method __init__ (line 477) | def __init__(self, dtype, device, checkpoint_path):
    method visual (line 498) | def visual(self, videos):

FILE: wan/modules/model.py
  function sinusoidal_embedding_1d (line 15) | def sinusoidal_embedding_1d(dim, position):
  function rope_params (line 29) | def rope_params(max_seq_len, dim, theta=10000):
  function rope_apply (line 40) | def rope_apply(x, grid_sizes, freqs):
  class WanRMSNorm (line 70) | class WanRMSNorm(nn.Module):
    method __init__ (line 72) | def __init__(self, dim, eps=1e-5):
    method forward (line 78) | def forward(self, x):
    method _norm (line 85) | def _norm(self, x):
  class WanLayerNorm (line 89) | class WanLayerNorm(nn.LayerNorm):
    method __init__ (line 91) | def __init__(self, dim, eps=1e-6, elementwise_affine=False):
    method forward (line 94) | def forward(self, x):
  class WanSelfAttention (line 102) | class WanSelfAttention(nn.Module):
    method __init__ (line 104) | def __init__(self,
    method forward (line 127) | def forward(self, x, seq_lens, grid_sizes, freqs):
  class WanT2VCrossAttention (line 159) | class WanT2VCrossAttention(WanSelfAttention):
    method forward (line 161) | def forward(self, x, context, context_lens, crossattn_cache=None):
  class WanGanCrossAttention (line 197) | class WanGanCrossAttention(WanSelfAttention):
    method forward (line 199) | def forward(self, x, context, crossattn_cache=None):
  class WanI2VCrossAttention (line 224) | class WanI2VCrossAttention(WanSelfAttention):
    method __init__ (line 226) | def __init__(self,
    method forward (line 240) | def forward(self, x, context, context_lens):
  class WanAttentionBlock (line 275) | class WanAttentionBlock(nn.Module):
    method __init__ (line 277) | def __init__(self,
    method forward (line 315) | def forward(
  class GanAttentionBlock (line 357) | class GanAttentionBlock(nn.Module):
    method __init__ (line 359) | def __init__(self,
    method forward (line 397) | def forward(
  class Head (line 439) | class Head(nn.Module):
    method __init__ (line 441) | def __init__(self, dim, out_dim, patch_size, eps=1e-6):
    method forward (line 456) | def forward(self, x, e):
  class MLPProj (line 469) | class MLPProj(torch.nn.Module):
    method __init__ (line 471) | def __init__(self, in_dim, out_dim):
    method forward (line 479) | def forward(self, image_embeds):
  class RegisterTokens (line 484) | class RegisterTokens(nn.Module):
    method __init__ (line 485) | def __init__(self, num_registers: int, dim: int):
    method forward (line 490) | def forward(self):
    method reset_parameters (line 493) | def reset_parameters(self):
  class WanModel (line 497) | class WanModel(ModelMixin, ConfigMixin):
    method __init__ (line 509) | def __init__(self,
    method _set_gradient_checkpointing (line 623) | def _set_gradient_checkpointing(self, module, value=False):
    method forward (line 626) | def forward(
    method _forward (line 637) | def _forward(
    method _forward_classify (line 773) | def _forward_classify(
    method unpatchify (line 876) | def unpatchify(self, x, grid_sizes, c=None):
    method init_weights (line 901) | def init_weights(self):

FILE: wan/modules/t5.py
  function fp16_clamp (line 20) | def fp16_clamp(x):
  function init_weights (line 27) | def init_weights(m):
  class GELU (line 46) | class GELU(nn.Module):
    method forward (line 48) | def forward(self, x):
  class T5LayerNorm (line 53) | class T5LayerNorm(nn.Module):
    method __init__ (line 55) | def __init__(self, dim, eps=1e-6):
    method forward (line 61) | def forward(self, x):
  class T5Attention (line 69) | class T5Attention(nn.Module):
    method __init__ (line 71) | def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
    method forward (line 86) | def forward(self, x, context=None, mask=None, pos_bias=None):
  class T5FeedForward (line 123) | class T5FeedForward(nn.Module):
    method __init__ (line 125) | def __init__(self, dim, dim_ffn, dropout=0.1):
    method forward (line 136) | def forward(self, x):
  class T5SelfAttention (line 144) | class T5SelfAttention(nn.Module):
    method __init__ (line 146) | def __init__(self,
    method forward (line 170) | def forward(self, x, mask=None, pos_bias=None):
  class T5CrossAttention (line 178) | class T5CrossAttention(nn.Module):
    method __init__ (line 180) | def __init__(self,
    method forward (line 206) | def forward(self,
  class T5RelativeEmbedding (line 221) | class T5RelativeEmbedding(nn.Module):
    method __init__ (line 223) | def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
    method forward (line 233) | def forward(self, lq, lk):
    method _relative_position_bucket (line 245) | def _relative_position_bucket(self, rel_pos):
  class T5Encoder (line 267) | class T5Encoder(nn.Module):
    method __init__ (line 269) | def __init__(self,
    method forward (line 303) | def forward(self, ids, mask=None):
  class T5Decoder (line 315) | class T5Decoder(nn.Module):
    method __init__ (line 317) | def __init__(self,
    method forward (line 351) | def forward(self, ids, mask=None, encoder_states=None, encoder_mask=No...
  class T5Model (line 372) | class T5Model(nn.Module):
    method __init__ (line 374) | def __init__(self,
    method forward (line 408) | def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
  function _t5 (line 415) | def _t5(name,
  function umt5_xxl (line 456) | def umt5_xxl(**kwargs):
  class T5EncoderModel (line 472) | class T5EncoderModel:
    method __init__ (line 474) | def __init__(
    method __call__ (line 506) | def __call__(self, texts, device):

FILE: wan/modules/tokenizers.py
  function basic_clean (line 12) | def basic_clean(text):
  function whitespace_clean (line 18) | def whitespace_clean(text):
  function canonicalize (line 24) | def canonicalize(text, keep_punctuation_exact_string=None):
  class HuggingfaceTokenizer (line 37) | class HuggingfaceTokenizer:
    method __init__ (line 39) | def __init__(self, name, seq_len=None, clean=None, **kwargs):
    method __call__ (line 49) | def __call__(self, sequence, **kwargs):
    method _clean (line 75) | def _clean(self, text):

FILE: wan/modules/vae.py
  class CausalConv3d (line 17) | class CausalConv3d(nn.Conv3d):
    method __init__ (line 22) | def __init__(self, *args, **kwargs):
    method forward (line 28) | def forward(self, x, cache_x=None):
  class RMS_norm (line 39) | class RMS_norm(nn.Module):
    method __init__ (line 41) | def __init__(self, dim, channel_first=True, images=True, bias=False):
    method forward (line 51) | def forward(self, x):
  class Upsample (line 57) | class Upsample(nn.Upsample):
    method forward (line 59) | def forward(self, x):
  class Resample (line 66) | class Resample(nn.Module):
    method __init__ (line 68) | def __init__(self, dim, mode):
    method forward (line 101) | def forward(self, x, feat_cache=None, feat_idx=[0]):
    method init_weight (line 162) | def init_weight(self, conv):
    method init_weight2 (line 174) | def init_weight2(self, conv):
  class ResidualBlock (line 186) | class ResidualBlock(nn.Module):
    method __init__ (line 188) | def __init__(self, in_dim, out_dim, dropout=0.0):
    method forward (line 202) | def forward(self, x, feat_cache=None, feat_idx=[0]):
  class AttentionBlock (line 223) | class AttentionBlock(nn.Module):
    method __init__ (line 228) | def __init__(self, dim):
    method forward (line 240) | def forward(self, x):
  class Encoder3d (line 265) | class Encoder3d(nn.Module):
    method __init__ (line 267) | def __init__(self,
    method forward (line 318) | def forward(self, x, feat_cache=None, feat_idx=[0]):
  class Decoder3d (line 369) | class Decoder3d(nn.Module):
    method __init__ (line 371) | def __init__(self,
    method forward (line 423) | def forward(self, x, feat_cache=None, feat_idx=[0]):
  function count_conv3d (line 475) | def count_conv3d(model):
  class WanVAE_ (line 483) | class WanVAE_(nn.Module):
    method __init__ (line 485) | def __init__(self,
    method forward (line 511) | def forward(self, x):
    method encode (line 517) | def encode(self, x, scale):
    method decode (line 545) | def decode(self, z, scale):
    method cached_decode (line 571) | def cached_decode(self, z, scale):
    method sample (line 595) | def sample(self, imgs, deterministic=False):
    method clear_cache (line 602) | def clear_cache(self):
  function _video_vae (line 612) | def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
  class WanVAE (line 639) | class WanVAE:
    method __init__ (line 641) | def __init__(self,
    method encode (line 667) | def encode(self, videos):
    method decode (line 677) | def decode(self, zs):

FILE: wan/modules/xlm_roberta.py
  class SelfAttention (line 10) | class SelfAttention(nn.Module):
    method __init__ (line 12) | def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5):
    method forward (line 27) | def forward(self, x, mask):
  class AttentionBlock (line 49) | class AttentionBlock(nn.Module):
    method __init__ (line 51) | def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5):
    method forward (line 66) | def forward(self, x, mask):
  class XLMRoberta (line 76) | class XLMRoberta(nn.Module):
    method __init__ (line 81) | def __init__(self,
    method forward (line 118) | def forward(self, ids):
  function xlm_roberta_large (line 146) | def xlm_roberta_large(pretrained=False,

FILE: wan/text2video.py
  class WanT2V (line 26) | class WanT2V:
    method __init__ (line 28) | def __init__(
    method generate (line 110) | def generate(self,

FILE: wan/utils/fm_solvers.py
  function get_sampling_sigmas (line 22) | def get_sampling_sigmas(sampling_steps, shift):
  function retrieve_timesteps (line 29) | def retrieve_timesteps(
  class FlowDPMSolverMultistepScheduler (line 69) | class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
    method __init__ (line 129) | def __init__(
    method step_index (line 202) | def step_index(self):
    method begin_index (line 209) | def begin_index(self):
    method set_begin_index (line 216) | def set_begin_index(self, begin_index: int = 0):
    method set_timesteps (line 226) | def set_timesteps(
    method _threshold_sample (line 292) | def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
    method _sigma_to_t (line 330) | def _sigma_to_t(self, sigma):
    method _sigma_to_alpha_sigma_t (line 333) | def _sigma_to_alpha_sigma_t(self, sigma):
    method time_shift (line 337) | def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
    method convert_model_output (line 341) | def convert_model_output(
    method dpm_solver_first_order_update (line 415) | def dpm_solver_first_order_update(
    method multistep_dpm_solver_second_order_update (line 486) | def multistep_dpm_solver_second_order_update(
    method multistep_dpm_solver_third_order_update (line 596) | def multistep_dpm_solver_third_order_update(
    method index_for_timestep (line 679) | def index_for_timestep(self, timestep, schedule_timesteps=None):
    method _init_step_index (line 693) | def _init_step_index(self, timestep):
    method step (line 706) | def step(
    method scale_model_input (line 800) | def scale_model_input(self, sample: torch.Tensor, *args,
    method add_noise (line 815) | def add_noise(
    method __len__ (line 856) | def __len__(self):

FILE: wan/utils/fm_solvers_unipc.py
  class FlowUniPCMultistepScheduler (line 20) | class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
    method __init__ (line 77) | def __init__(
    method step_index (line 135) | def step_index(self):
    method begin_index (line 142) | def begin_index(self):
    method set_begin_index (line 149) | def set_begin_index(self, begin_index: int = 0):
    method set_timesteps (line 160) | def set_timesteps(
    method _threshold_sample (line 230) | def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
    method _sigma_to_t (line 269) | def _sigma_to_t(self, sigma):
    method _sigma_to_alpha_sigma_t (line 272) | def _sigma_to_alpha_sigma_t(self, sigma):
    method time_shift (line 276) | def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
    method convert_model_output (line 279) | def convert_model_output(
    method multistep_uni_p_bh_update (line 350) | def multistep_uni_p_bh_update(
    method multistep_uni_c_bh_update (line 486) | def multistep_uni_c_bh_update(
    method index_for_timestep (line 628) | def index_for_timestep(self, timestep, schedule_timesteps=None):
    method _init_step_index (line 643) | def _init_step_index(self, timestep):
    method step (line 655) | def step(self,
    method scale_model_input (line 741) | def scale_model_input(self, sample: torch.Tensor, *args,
    method add_noise (line 758) | def add_noise(
    method __len__ (line 799) | def __len__(self):

FILE: wan/utils/prompt_extend.py
  class PromptOutput (line 101) | class PromptOutput(object):
    method add_custom_field (line 108) | def add_custom_field(self, key: str, value) -> None:
  class PromptExpander (line 112) | class PromptExpander:
    method __init__ (line 114) | def __init__(self, model_name, is_vl=False, device=0, **kwargs):
    method extend_with_img (line 119) | def extend_with_img(self,
    method extend (line 128) | def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
    method decide_system_prompt (line 131) | def decide_system_prompt(self, tar_lang="ch"):
    method __call__ (line 138) | def __call__(self,
  class DashScopePromptExpander (line 157) | class DashScopePromptExpander(PromptExpander):
    method __init__ (line 159) | def __init__(self,
    method extend (line 196) | def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
    method extend_with_img (line 232) | def extend_with_img(self,
  class QwenPromptExpander (line 300) | class QwenPromptExpander(PromptExpander):
    method __init__ (line 309) | def __init__(self, model_name=None, device=0, is_vl=False, **kwargs):
    method extend (line 366) | def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
    method extend_with_img (line 397) | def extend_with_img(self,

FILE: wan/utils/qwen_vl_utils.py
  function round_by_factor (line 39) | def round_by_factor(number: int, factor: int) -> int:
  function ceil_by_factor (line 44) | def ceil_by_factor(number: int, factor: int) -> int:
  function floor_by_factor (line 49) | def floor_by_factor(number: int, factor: int) -> int:
  function smart_resize (line 54) | def smart_resize(height: int,
  function fetch_image (line 85) | def fetch_image(ele: dict[str, str | Image.Image],
  function smart_nframes (line 133) | def smart_nframes(
  function _read_video_torchvision (line 177) | def _read_video_torchvision(ele: dict,) -> torch.Tensor:
  function is_decord_available (line 215) | def is_decord_available() -> bool:
  function _read_video_decord (line 221) | def _read_video_decord(ele: dict,) -> torch.Tensor:
  function get_video_reader_backend (line 261) | def get_video_reader_backend() -> str:
  function fetch_video (line 274) | def fetch_video(
  function extract_vision_info (line 328) | def extract_vision_info(
  function process_vision_info (line 344) | def process_vision_info(

FILE: wan/utils/utils.py
  function rand_name (line 14) | def rand_name(length=8, suffix=''):
  function cache_video (line 23) | def cache_video(tensor,
  function cache_image (line 64) | def cache_image(tensor,
  function str2bool (line 94) | def str2bool(v):