SYMBOL INDEX (1104 symbols across 103 files)

FILE: demo.py
  function initialize_vae_decoder (line 60) | def initialize_vae_decoder(use_taehv=False, use_trt=False):
  function tensor_to_base64_frame (line 162) | def tensor_to_base64_frame(frame_tensor):
  function frame_sender_worker (line 190) | def frame_sender_worker():
  function generate_video_stream (line 242) | def generate_video_stream(prompt, seed, enable_torch_compile=False, enab...
  function generate_mp4_from_images (line 509) | def generate_mp4_from_images(image_directory, output_video_path, fps=24):
  function calculate_sha256 (line 533) | def calculate_sha256(data):
  function handle_connect (line 543) | def handle_connect():
  function handle_disconnect (line 549) | def handle_disconnect():
  function handle_start_generation (line 554) | def handle_start_generation(data):
  function handle_stop_generation (line 597) | def handle_stop_generation():
  function index (line 614) | def index():
  function api_status (line 619) | 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: get_causal_ode_data_chunkwise.py
  function init_model (line 13) | def init_model(device):
  function main (line 32) | def main():

FILE: get_causal_ode_data_framewise.py
  function init_model (line 13) | def init_model(device):
  function main (line 32) | def main():

FILE: get_causal_ode_data_kv_optimized.py
  function normalize_generator_state_dict (line 31) | def normalize_generator_state_dict(state_dict: dict) -> dict:
  function init_model (line 47) | def init_model(
  function prepare_clean_latent (line 73) | def prepare_clean_latent(
  function main (line 92) | def main():

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

FILE: long_video/app.py
  function _ensure_gpu (line 22) | def _ensure_gpu():
  function _load_pipeline (line 29) | def _load_pipeline(config_path: str, checkpoint_path: Optional[str], use...
  function build_predict (line 78) | def build_predict(config_path: str, checkpoint_path: Optional[str], outp...
  function main (line 120) | def main():

FILE: long_video/inference.py
  function remove_fsdp_prefix (line 70) | def remove_fsdp_prefix(state_dict):
  function encode (line 114) | def encode(self, videos: torch.Tensor) -> torch.Tensor:

FILE: long_video/model/base.py
  class BaseModel (line 12) | class BaseModel(nn.Module):
    method __init__ (line 13) | def __init__(self, args, device):
    method _initialize_models (line 26) | def _initialize_models(self, args, device):
    method _get_timestep (line 48) | def _get_timestep(
  class RollingForcingModel (line 98) | class RollingForcingModel(BaseModel):
    method __init__ (line 99) | def __init__(self, args, device):
    method _run_generator (line 103) | def _run_generator(
    method _consistency_backward_simulation (line 182) | def _consistency_backward_simulation(
    method _initialize_inference_pipeline (line 214) | def _initialize_inference_pipeline(self):

FILE: long_video/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: long_video/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: long_video/model/dmd.py
  class DMD (line 9) | class DMD(RollingForcingModel):
    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 128) | def compute_distribution_matching_loss(
    method generator_loss (line 196) | def generator_loss(
    method critic_loss (line 237) | def critic_loss(

FILE: long_video/model/gan.py
  class GAN (line 10) | class GAN(RollingForcingModel):
    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: long_video/model/ode_regression.py
  class ODERegression (line 9) | class ODERegression(BaseModel):
    method __init__ (line 10) | def __init__(self, args, device):
    method _initialize_models (line 45) | def _initialize_models(self, args, device):
    method _prepare_generator_input (line 59) | def _prepare_generator_input(self, ode_latent: torch.Tensor, tf=False,...
    method generator_loss (line 95) | def generator_loss(self, ode_latent: torch.Tensor, conditional_dict: d...

FILE: long_video/model/sid.py
  class SiD (line 8) | class SiD(RollingForcingModel):
    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: long_video/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: long_video/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: long_video/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: long_video/pipeline/rolling_forcing_inference.py
  class CausalInferencePipeline (line 7) | class CausalInferencePipeline(torch.nn.Module):
    method __init__ (line 8) | def __init__(
    method inference_rolling_forcing (line 45) | def inference_rolling_forcing(
    method _initialize_kv_cache (line 338) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 360) | def _initialize_crossattn_cache(self, batch_size, dtype, device):

FILE: long_video/pipeline/rolling_forcing_training.py
  class RollingForcingTrainingPipeline (line 8) | class RollingForcingTrainingPipeline:
    method __init__ (line 9) | def __init__(self,
    method generate_and_sync_list (line 41) | def generate_and_sync_list(self, num_blocks, num_denoising_steps, devi...
    method generate_list (line 60) | def generate_list(self, num_blocks, num_denoising_steps, device):
    method inference_with_rolling_forcing (line 75) | def inference_with_rolling_forcing(
    method inference_with_self_forcing (line 256) | def inference_with_self_forcing(
    method _initialize_kv_cache (line 436) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 452) | def _initialize_crossattn_cache(self, batch_size, dtype, device):

FILE: long_video/train.py
  function main (line 8) | def main():

FILE: long_video/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: long_video/trainer/distillation.py
  class Trainer (line 20) | class Trainer:
    method __init__ (line 21) | def __init__(self, config):
    method save (line 176) | def save(self):
    method fwdbwd_one_step (line 203) | def fwdbwd_one_step(self, batch, train_generator):
    method generate_video (line 276) | def generate_video(self, pipeline, prompts, image=None):
    method train (line 306) | def train(self):

FILE: long_video/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: long_video/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: long_video/utils/dataset.py
  class TextDataset (line 12) | class TextDataset(Dataset):
    method __init__ (line 13) | def __init__(self, prompt_path, extended_prompt_path=None):
    method __len__ (line 24) | def __len__(self):
    method __getitem__ (line 27) | def __getitem__(self, idx):
  class ODERegressionLMDBDataset (line 37) | class ODERegressionLMDBDataset(Dataset):
    method __init__ (line 38) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 45) | def __len__(self):
    method __getitem__ (line 48) | def __getitem__(self, idx):
  class ShardingLMDBDataset (line 72) | class ShardingLMDBDataset(Dataset):
    method __init__ (line 73) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 96) | def __len__(self):
    method __getitem__ (line 99) | def __getitem__(self, idx):
  class TextImagePairDataset (line 127) | class TextImagePairDataset(Dataset):
    method __init__ (line 128) | def __init__(
    method __len__ (line 182) | def __len__(self):
    method __getitem__ (line 185) | def __getitem__(self, idx):
  function cycle (line 217) | def cycle(dl):

FILE: long_video/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 70) | def barrier():
  function launch_distributed_job (line 75) | def launch_distributed_job(backend: str = "nccl"):
  class EMA_FSDP (line 91) | class EMA_FSDP:
    method __init__ (line 92) | def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
    method _init_shadow (line 98) | def _init_shadow(self, fsdp_module):
    method update (line 105) | def update(self, fsdp_module):
    method state_dict (line 113) | def state_dict(self):
    method load_state_dict (line 116) | def load_state_dict(self, sd):
    method copy_to (line 119) | def copy_to(self, fsdp_module):

FILE: long_video/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 56) | def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape...

FILE: long_video/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: long_video/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: long_video/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: long_video/utils/wan_wrapper.py
  class WanTextEncoder (line 14) | class WanTextEncoder(torch.nn.Module):
    method __init__ (line 15) | def __init__(self) -> None:
    method device (line 33) | def device(self):
    method forward (line 37) | def forward(self, text_prompts: List[str]) -> dict:
  class WanVAEWrapper (line 53) | class WanVAEWrapper(torch.nn.Module):
    method __init__ (line 54) | def __init__(self):
    method encode_to_latent (line 73) | def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
    method decode_to_pixel (line 89) | def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = Fals...
  class WanDiffusionWrapper (line 115) | class WanDiffusionWrapper(torch.nn.Module):
    method __init__ (line 116) | def __init__(
    method enable_gradient_checkpointing (line 144) | def enable_gradient_checkpointing(self) -> None:
    method adding_cls_branch (line 147) | def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_di...
    method _convert_flow_pred_to_x0 (line 169) | def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch....
    method _convert_x0_to_flow_pred (line 196) | def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: tor...
    method forward (line 218) | def forward(
    method get_scheduler (line 293) | def get_scheduler(self) -> SchedulerInterface:
    method post_init (line 307) | def post_init(self):

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

FILE: long_video/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: long_video/wan/image2video.py
  class WanI2V (line 29) | class WanI2V:
    method __init__ (line 31) | def __init__(
    method generate (line 129) | def generate(self,

FILE: long_video/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: long_video/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 87) | def forward(
  class CausalWanAttentionBlock (line 309) | class CausalWanAttentionBlock(nn.Module):
    method __init__ (line 311) | def __init__(self,
    method forward (line 349) | def forward(
  class CausalHead (line 405) | class CausalHead(nn.Module):
    method __init__ (line 407) | def __init__(self, dim, out_dim, patch_size, eps=1e-6):
    method forward (line 422) | def forward(self, x, e):
  class CausalWanModel (line 436) | class CausalWanModel(ModelMixin, ConfigMixin):
    method __init__ (line 448) | def __init__(self,
    method _set_gradient_checkpointing (line 569) | def _set_gradient_checkpointing(self, module, value=False):
    method _prepare_blockwise_causal_attn_mask (line 573) | def _prepare_blockwise_causal_attn_mask(
    method _prepare_teacher_forcing_mask (line 631) | def _prepare_teacher_forcing_mask(
    method _prepare_blockwise_causal_attn_mask_i2v (line 719) | def _prepare_blockwise_causal_attn_mask_i2v(
    method _forward_inference (line 779) | def _forward_inference(
    method _forward_train (line 912) | def _forward_train(
    method forward (line 1070) | def forward(
    method unpatchify (line 1080) | def unpatchify(self, x, grid_sizes):
    method init_weights (line 1105) | def init_weights(self):

FILE: long_video/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 53) | class SelfAttention(nn.Module):
    method __init__ (line 55) | def __init__(self,
    method forward (line 74) | def forward(self, x):
  class SwiGLU (line 94) | class SwiGLU(nn.Module):
    method __init__ (line 96) | def __init__(self, dim, mid_dim):
    method forward (line 106) | def forward(self, x):
  class AttentionBlock (line 112) | class AttentionBlock(nn.Module):
    method __init__ (line 114) | def __init__(self,
    method forward (line 146) | def forward(self, x):
  class AttentionPool (line 156) | class AttentionPool(nn.Module):
    method __init__ (line 158) | def __init__(self,
    method forward (line 186) | def forward(self, x):
  class VisionTransformer (line 209) | class VisionTransformer(nn.Module):
    method __init__ (line 211) | def __init__(self,
    method forward (line 279) | def forward(self, x, interpolation=False, use_31_block=False):
  class XLMRobertaWithHead (line 303) | class XLMRobertaWithHead(XLMRoberta):
    method __init__ (line 305) | def __init__(self, **kwargs):
    method forward (line 315) | def forward(self, ids):
  class XLMRobertaCLIP (line 328) | class XLMRobertaCLIP(nn.Module):
    method __init__ (line 330) | def __init__(self,
    method forward (line 406) | def forward(self, imgs, txt_ids):
    method param_groups (line 418) | def param_groups(self):
  function _clip (line 434) | def _clip(pretrained=False,
  function clip_xlm_roberta_vit_h_14 (line 471) | def clip_xlm_roberta_vit_h_14(
  class CLIPModel (line 501) | class CLIPModel:
    method __init__ (line 503) | def __init__(self, dtype, device, checkpoint_path, tokenizer_path):
    method visual (line 527) | def visual(self, videos):

FILE: long_video/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: long_video/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: long_video/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: long_video/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: long_video/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: long_video/wan/text2video.py
  class WanT2V (line 26) | class WanT2V:
    method __init__ (line 28) | def __init__(
    method generate (line 110) | def generate(self,

FILE: long_video/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: long_video/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: long_video/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: long_video/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: long_video/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):

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 40) | def _initialize_models(self, args, device):
    method _get_timestep (line 63) | def _get_timestep(
  class SelfForcingModel (line 113) | class SelfForcingModel(BaseModel):
    method __init__ (line 114) | def __init__(self, args, device):
    method _run_generator (line 118) | def _run_generator(
    method _consistency_backward_simulation (line 205) | def _consistency_backward_simulation(
    method _initialize_inference_pipeline (line 229) | def _initialize_inference_pipeline(self):
  class TeacherForcingModel (line 251) | class TeacherForcingModel(BaseModel):
    method __init__ (line 252) | def __init__(self, args, device):
    method _run_generator (line 256) | def _run_generator(
    method _consistency_backward_simulation_tf (line 344) | def _consistency_backward_simulation_tf(
    method _initialize_inference_pipeline_tf (line 372) | def _initialize_inference_pipeline_tf(self):
  class BidirectionalModel (line 391) | class BidirectionalModel(BaseModel):
    method __init__ (line 392) | def __init__(self, args, device):
    method _run_generator (line 396) | def _run_generator(
    method _consistency_backward_simulation_bidirectional (line 476) | def _consistency_backward_simulation_bidirectional(
    method _initialize_inference_pipeline_bidirectional (line 502) | def _initialize_inference_pipeline_bidirectional(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 256) | def generator_loss(
    method critic_loss (line 297) | 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 35) | def _initialize_models(self, args, device):
    method generator_loss (line 48) | 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 56) | def _compute_kl_grad(
    method compute_distribution_matching_loss (line 130) | def compute_distribution_matching_loss(
    method generator_loss (line 199) | def generator_loss(
    method critic_loss (line 240) | def critic_loss(
    method _prepare_generator_input (line 338) | def _prepare_generator_input(self, ode_latent: torch.Tensor, tf=False,...

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/naive_consistency.py
  class NaiveConsistency (line 9) | class NaiveConsistency(BaseModel):
    method __init__ (line 10) | def __init__(self, args, device):
    method _initialize_models (line 66) | def _initialize_models(self, args, device):
    method generator_loss (line 85) | def generator_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 45) | def _initialize_models(self, args, device):
    method _prepare_generator_input (line 59) | def _prepare_generator_input(self, ode_latent: torch.Tensor, tf=False,...
    method generator_loss (line 95) | 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 8) | class BidirectionalTrainingPipeline:
    method __init__ (line 9) | def __init__(self,
    method generate_and_sync_list (line 44) | def generate_and_sync_list(self, num_blocks, num_denoising_steps, devi...
    method inference_with_trajectory (line 64) | def inference_with_trajectory(

FILE: pipeline/causal_diffusion_inference.py
  class CausalDiffusionInferencePipeline (line 10) | class CausalDiffusionInferencePipeline(torch.nn.Module):
    method __init__ (line 11) | def __init__(
    method inference (line 51) | def inference(
    method inference_for_cd (line 277) | def inference_for_cd(
    method inference_for_genuine_cd (line 474) | def inference_for_genuine_cd(
    method _initialize_kv_cache (line 605) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 635) | def _initialize_crossattn_cache(self, batch_size, dtype, device):
    method _initialize_sample_scheduler (line 656) | def _initialize_sample_scheduler(self, noise, sampling_steps=-1):

FILE: pipeline/causal_inference.py
  class CausalInferencePipeline (line 10) | class CausalInferencePipeline(torch.nn.Module):
    method __init__ (line 11) | def __init__(
    method inference (line 65) | def inference(
    method _initialize_kv_cache (line 337) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 359) | 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 48) | def generate_and_sync_list(self, num_blocks, num_denoising_steps, devi...
    method inference_with_trajectory (line 68) | def inference_with_trajectory(
    method _initialize_kv_cache (line 288) | def _initialize_kv_cache(self, batch_size, dtype, device):
    method _initialize_crossattn_cache (line 304) | def _initialize_crossattn_cache(self, batch_size, dtype, device):

FILE: pipeline/teacher_forcing_training.py
  class TeacherForcingTrainingPipeline (line 8) | class TeacherForcingTrainingPipeline:
    method __init__ (line 9) | def __init__(self,
    method generate_and_sync_list (line 44) | def generate_and_sync_list(self, num_blocks, num_denoising_steps, devi...
    method inference_with_trajectory (line 64) | def inference_with_trajectory(

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

FILE: trainer/diffusion.py
  class Trainer (line 20) | class Trainer:
    method __init__ (line 21) | def __init__(self, config):
    method save (line 168) | def save(self):
    method train_one_step (line 191) | def train_one_step(self, batch):
    method train (line 261) | def train(self):

FILE: trainer/distillation.py
  class Trainer (line 16) | class Trainer:
    method __init__ (line 17) | def __init__(self, config):
    method save (line 188) | def save(self):
    method save_critic (line 212) | def save_critic(self):
    method fwdbwd_one_step (line 229) | def fwdbwd_one_step(self, batch, train_generator, clean_latent=None):
    method train (line 303) | 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/naive_cd.py
  class Trainer (line 20) | class Trainer:
    method __init__ (line 21) | def __init__(self, config):
    method save (line 183) | def save(self):
    method fwdbwd_one_step (line 206) | def fwdbwd_one_step(self, batch, clean_latent=None):
    method train (line 251) | def train(self):

FILE: trainer/ode.py
  class Trainer (line 19) | class Trainer:
    method __init__ (line 20) | def __init__(self, config):
    method save (line 132) | def save(self):
    method train_one_step (line 148) | def train_one_step(self, loss_scale=1.0):
    method train (line 217) | def train(self):

FILE: utils/create_lmdb_iterative.py
  function store_arrays_to_lmdb (line 10) | def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
  function get_array_shape_from_lmdb (line 27) | def get_array_shape_from_lmdb(env, array_name):
  function process_data_dict (line 34) | def process_data_dict(data_dict, seen_prompts):
  function retrieve_row_from_lmdb (line 59) | def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape...
  function main (line 78) | def main():

FILE: utils/dataset.py
  class TextDataset (line 12) | class TextDataset(Dataset):
    method __init__ (line 13) | def __init__(self, prompt_path, extended_prompt_path=None):
    method __len__ (line 24) | def __len__(self):
    method __getitem__ (line 27) | def __getitem__(self, idx):
  class ODERegressionLMDBDataset (line 37) | class ODERegressionLMDBDataset(Dataset):
    method __init__ (line 38) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 45) | def __len__(self):
    method __getitem__ (line 48) | def __getitem__(self, idx):
  class LatentLMDBDataset (line 75) | class LatentLMDBDataset(Dataset):
    method __init__ (line 76) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 83) | def __len__(self):
    method __getitem__ (line 86) | def __getitem__(self, idx):
  class ShardingLMDBDataset (line 110) | class ShardingLMDBDataset(Dataset):
    method __init__ (line 111) | def __init__(self, data_path: str, max_pair: int = int(1e8)):
    method __len__ (line 134) | def __len__(self):
    method __getitem__ (line 137) | def __getitem__(self, idx):
  class TextImagePairDataset (line 166) | class TextImagePairDataset(Dataset):
    method __init__ (line 167) | def __init__(
    method __len__ (line 221) | def __len__(self):
    method __getitem__ (line 224) | def __getitem__(self, idx):
  function cycle (line 257) | 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 70) | def barrier():
  function launch_distributed_job (line 75) | def launch_distributed_job(backend: str = "nccl"):
  class EMA_FSDP (line 91) | class EMA_FSDP:
    method __init__ (line 92) | def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
    method _init_shadow (line 98) | def _init_shadow(self, fsdp_module):
    method update (line 103) | def update(self, fsdp_module):
    method state_dict (line 109) | def state_dict(self):
    method load_state_dict (line 112) | def load_state_dict(self, sd):
    method copy_to (line 115) | def copy_to(self, fsdp_module):
    method full_state_dict (line 121) | def full_state_dict(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 56) | 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/merge_and_get_clean.py
  function read_shape (line 8) | def read_shape(env, name):
  function list_array_names (line 14) | def list_array_names(env):
  function ensure_empty_dir (line 24) | def ensure_empty_dir(path):
  function safe_mapsize (line 28) | def safe_mapsize(env):
  function get_bytes (line 32) | def get_bytes(txn, key):
  function latents_bytes_to_out (line 37) | def latents_bytes_to_out(row_bytes, in_row_shape, out_row_shape):
  function merge_many (line 50) | def merge_many(src_dirs_all, dst_dir):
  function rm_dirs (line 140) | def rm_dirs(dirs, desc="remove dirs"):

FILE: utils/merge_lmdb.py
  function read_shape (line 8) | def read_shape(env, name):
  function list_array_names (line 15) | def list_array_names(env):
  function ensure_empty_dir (line 26) | def ensure_empty_dir(path):
  function safe_mapsize (line 31) | def safe_mapsize(env):
  function get_bytes (line 35) | def get_bytes(txn, key):
  function latents_bytes_to_out (line 41) | def latents_bytes_to_out(row_bytes, in_row_shape, out_row_shape):
  function merge_many (line 47) | def merge_many(src_dirs_all, dst_dir):
  function rm_dirs (line 142) | def rm_dirs(dirs, desc="remove dirs"):

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/ode_generation.py
  function merge_cfg_prompt_embeds (line 6) | def merge_cfg_prompt_embeds(
  function normalize_trajectory_indices (line 21) | def normalize_trajectory_indices(
  class CausalODETrajectoryGenerator (line 37) | class CausalODETrajectoryGenerator:
    method __init__ (line 38) | def __init__(
    method _make_kv_cache (line 55) | def _make_kv_cache(self, batch_size: int, device: torch.device) -> lis...
    method _make_crossattn_cache (line 81) | def _make_crossattn_cache(self, batch_size: int, device: torch.device)...
    method _batched_cfg_step (line 101) | def _batched_cfg_step(
    method _update_clean_cache (line 132) | def _update_clean_cache(
    method _generate_full (line 156) | def _generate_full(
    method _generate_blockwise_kv (line 196) | def _generate_blockwise_kv(
    method _assemble_selected_trajectory (line 267) | def _assemble_selected_trajectory(
    method generate (line 290) | def generate(

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 14) | class WanTextEncoder(torch.nn.Module):
    method __init__ (line 15) | def __init__(self) -> None:
    method device (line 33) | def device(self):
    method forward (line 37) | def forward(self, text_prompts: List[str]) -> dict:
  class WanVAEWrapper (line 53) | class WanVAEWrapper(torch.nn.Module):
    method __init__ (line 54) | def __init__(self):
    method encode_to_latent (line 73) | def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
    method decode_to_pixel (line 89) | def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = Fals...
  class WanDiffusionWrapper (line 115) | class WanDiffusionWrapper(torch.nn.Module):
    method __init__ (line 116) | def __init__(
    method enable_gradient_checkpointing (line 144) | def enable_gradient_checkpointing(self) -> None:
    method adding_cls_branch (line 147) | def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_di...
    method _convert_flow_pred_to_x0 (line 169) | def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch....
    method _convert_x0_to_flow_pred (line 196) | def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: tor...
    method forward (line 218) | def forward(
    method get_scheduler (line 294) | def get_scheduler(self) -> SchedulerInterface:
    method post_init (line 308) | 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 31) | def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
  class CausalWanSelfAttention (line 62) | class CausalWanSelfAttention(nn.Module):
    method __init__ (line 64) | def __init__(self,
    method forward (line 90) | def forward(
  class CausalWanAttentionBlock (line 248) | class CausalWanAttentionBlock(nn.Module):
    method __init__ (line 250) | def __init__(self,
    method forward (line 288) | def forward(
  class CausalHead (line 343) | class CausalHead(nn.Module):
    method __init__ (line 345) | def __init__(self, dim, out_dim, patch_size, eps=1e-6):
    method forward (line 360) | def forward(self, x, e):
  class CausalWanModel (line 374) | class CausalWanModel(ModelMixin, ConfigMixin):
    method __init__ (line 386) | def __init__(self,
    method _set_gradient_checkpointing (line 507) | def _set_gradient_checkpointing(self, module, value=False):
    method _prepare_blockwise_causal_attn_mask (line 511) | def _prepare_blockwise_causal_attn_mask(
    method _prepare_teacher_forcing_mask (line 569) | def _prepare_teacher_forcing_mask(
    method _prepare_blockwise_causal_attn_mask_i2v (line 657) | def _prepare_blockwise_causal_attn_mask_i2v(
    method _forward_inference (line 717) | def _forward_inference(
    method _forward_train (line 848) | def _forward_train(
    method forward (line 1007) | def forward(
    method unpatchify (line 1018) | def unpatchify(self, x, grid_sizes):
    method init_weights (line 1043) | 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 53) | class SelfAttention(nn.Module):
    method __init__ (line 55) | def __init__(self,
    method forward (line 74) | def forward(self, x):
  class SwiGLU (line 94) | class SwiGLU(nn.Module):
    method __init__ (line 96) | def __init__(self, dim, mid_dim):
    method forward (line 106) | def forward(self, x):
  class AttentionBlock (line 112) | class AttentionBlock(nn.Module):
    method __init__ (line 114) | def __init__(self,
    method forward (line 146) | def forward(self, x):
  class AttentionPool (line 156) | class AttentionPool(nn.Module):
    method __init__ (line 158) | def __init__(self,
    method forward (line 186) | def forward(self, x):
  class VisionTransformer (line 209) | class VisionTransformer(nn.Module):
    method __init__ (line 211) | def __init__(self,
    method forward (line 279) | def forward(self, x, interpolation=False, use_31_block=False):
  class XLMRobertaWithHead (line 303) | class XLMRobertaWithHead(XLMRoberta):
    method __init__ (line 305) | def __init__(self, **kwargs):
    method forward (line 315) | def forward(self, ids):
  class XLMRobertaCLIP (line 328) | class XLMRobertaCLIP(nn.Module):
    method __init__ (line 330) | def __init__(self,
    method forward (line 406) | def forward(self, imgs, txt_ids):
    method param_groups (line 418) | def param_groups(self):
  function _clip (line 434) | def _clip(pretrained=False,
  function clip_xlm_roberta_vit_h_14 (line 471) | def clip_xlm_roberta_vit_h_14(
  class CLIPModel (line 501) | class CLIPModel:
    method __init__ (line 503) | def __init__(self, dtype, device, checkpoint_path, tokenizer_path):
    method visual (line 527) | 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):