Functions — github.com/MatrixTeam-AI/RAIN

hub / github.com/MatrixTeam-AI/RAIN / functions

Functions352 in github.com/MatrixTeam-AI/RAIN

⨍ Functions 352◇ Types & classes 79

↓ 141 callersMethodto

(self, device)

src/scheduler/scheduler_lcm.py:254

↓ 40 callersMethodhead_to_batch_dim

r""" Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is the number of hea

src/models/attention_processor.py:547

↓ 16 callersFunctionconv

(n_in, n_out, **kwargs)

src/taesdv/taesdv.py:12

↓ 13 callersMethodnorm_encoder_hidden_states

r""" Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the `Attention` class.

src/models/attention_processor.py:664

↓ 13 callersMethodprepare_attention_mask

r""" Prepare the attention mask for the attention computation. Args: attention_mask (`torch.Tensor`): The

src/models/attention_processor.py:617

↓ 13 callersFunctionset_params

(name)

gradio_app.py:132

↓ 12 callersFunctiontorch_dfs

(model: torch.nn.Module)

src/models/mutual_self_attention.py:12

↓ 10 callersMethod__init__

( self, query_dim: int, cross_attention_dim: Optional[int] = None, heads: int

src/models/attention_processor.py:89

↓ 10 callersMethodbatch_to_head_dim

r""" Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads` is the number of h

src/models/attention_processor.py:530

↓ 10 callersMethodset_processor

r""" Set the attention processor to use. Args: processor (`AttnProcessor`): The attention processor to us

src/models/attention_processor.py:371

↓ 7 callersMethodget_attention_scores

r""" Compute the attention scores. Args: query (`torch.Tensor`): The query tensor. key (`torch.Tensor`): The

src/models/attention_processor.py:570

↓ 7 callersMethodupdate

(self, writer, dtype=torch.float16)

src/models/mutual_self_attention.py:359

↓ 6 callersMethod__init__

( self, in_channels: int, prev_output_channel: int, out_channels: int,

src/models/unet_2d_blocks.py:963

↓ 6 callersMethodclear

(self)

src/models/mutual_self_attention.py:409

↓ 6 callersMethodwrite

(self, frame_tensor)

src/taesdv/taesdv.py:190

↓ 5 callersMethoddetect

(self, face_image)

src/morpher.py:604

↓ 5 callersFunctionget_motion_module

(in_channels, motion_module_type: str, motion_module_kwargs: dict)

src/models/motion_module.py:34

↓ 5 callersMethodset_timesteps

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`,

src/scheduler/scheduler_lcm.py:316

↓ 4 callersMethod__init__

( self, in_channels: int, prev_output_channel: int, out_channels: int,

src/models/unet_3d_blocks.py:805

↓ 4 callersMethod__init__

(self, d_model, dropout=0.0, max_len=24)

src/models/motion_module.py:263

↓ 4 callersFunctionexport_onnx

( model, onnx_path: str, opt_image_height: int, opt_image_width: int, opt_batch_size: int,

src/modeling/onnx_export.py:27

↓ 4 callersFunctionmatch_dtype

(a, b)

src/modeling/engine_model.py:48

↓ 4 callersFunctionmatch_shape

(a, b)

src/modeling/engine_model.py:37

↓ 4 callersMethodprefill

(self, skip_check=SKIP_ENGINE_MODEL_CHECK, **inputs)

src/modeling/engine_model.py:163

↓ 4 callersMethodstep

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned m

src/scheduler/scheduler_lcm.py:391

↓ 4 callersMethodverify

(self, img, size=None, centercrop=False)

src/morpher.py:722

↓ 3 callersMethod__init__

(self, pose_guider, unet, vae, scheduler, timestep, w_embedding)

src/modeling/framed_models.py:6

↓ 3 callersMethod_reshape

(hid_states, h)

src/models/attention_processor.py:2244

↓ 3 callersMethodclear_queue

(self)

src/morpher.py:450

↓ 3 callersMethodfrom_pretrained_2d

( cls, pretrained_model_path: PathLike, motion_module_path=None, subfolder=Non

src/models/unet_3d.py:615

↓ 3 callersFunctionget_down_block

( down_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channel

src/models/unet_2d_blocks.py:20

↓ 3 callersFunctionget_up_block

( up_block_type: str, num_layers: int, in_channels: int, out_channels: int, prev_output_ch

src/models/unet_2d_blocks.py:102

↓ 3 callersMethodpush_input

(self, x, mix=1.0, eager_mode = False, buffer_limit=6)

src/morpher.py:626

↓ 3 callersMethodset_translations

(self, translations)

src/morpher.py:617

↓ 2 callersMethod__init__

(self, n_in, n_out)

src/taesdv/taesdv.py:20

↓ 2 callersMethod__init__

( self, channels, use_conv=False, use_conv_transpose=False, out_channe

src/models/resnet.py:32

↓ 2 callersFunction_get_3rd_point

To calculate the affine matrix, three pairs of points are required. This function is used to get the 3rd point, given 2D points a & b. The 3r

src/dwpose/onnxpose.py:183

↓ 2 callersFunctionalpha_bar_fn

(t)

src/scheduler/scheduler_lcm.py:79

↓ 2 callersMethoddecode_video

Decode a sequence of frames. Args: x: input NTCHW latent (C=4) tensor with ~Gaussian values. parallel: if true, all f

src/taesdv/taesdv.py:110

↓ 2 callersMethodencode_frame

Encode a single RGB timestep to latents. Args: x: input NCHW RGB (C=3) tensor with values in [0, 1]. Returns NCHW latent

src/taesdv/taesdv.py:58

↓ 2 callersMethodencode_video

Encode a sequence of frames. Args: x: input NTCHW RGB (C=3) tensor with values in [0, 1]. parallel: if True, all fram

src/taesdv/taesdv.py:91

↓ 2 callersMethodfetch_one_frame

(self, wait_after=0.0)

src/morpher.py:607

↓ 2 callersMethodforward

(self, hidden_states)

src/models/resnet.py:255

↓ 2 callersMethodkick_start

(self, wait_after=0.01)

src/morpher.py:642

↓ 2 callersFunctionparam2translations

(face_length_ratio=0.75, eye_height_ratio=3.0, eye_width_ratio=1.6, mouth_height_ratio=1.2, mouth_width_ratio=

gradio_app.py:208

↓ 2 callersMethodset_patterns

(self, patterns)

src/morpher.py:620

↓ 2 callersMethodstop

(self)

src/morpher.py:683

↓ 2 callersFunctionzero_module

(module)

src/models/motion_module.py:15

↓ 1 callersMethod__init__

( self, dim: int, num_attention_heads: int, attention_head_dim: int, d

src/models/attention.py:48

↓ 1 callersMethod_encode_prompt

( self, prompt, device, num_videos_per_prompt, do_classifier_free_guid

src/pipeline/pipeline_pose2vid_lcm.py:186

↓ 1 callersFunction_fix_aspect_ratio

Extend the scale to match the given aspect ratio. Args: scale (np.ndarray): The image scale (w, h) in shape (2, ) aspect_ratio (f

src/dwpose/onnxpose.py:150

↓ 1 callersFunction_rotate_point

Rotate a point by an angle. Args: pt (np.ndarray): 2D point coordinates (x, y) in shape (2, ) angle_rad (float): rotation angle i

src/dwpose/onnxpose.py:168

↓ 1 callersMethod_threshold_sample

"Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the prediction of x_0 at t

src/scheduler/scheduler_lcm.py:283

↓ 1 callersFunctionauto_cast_manager

(enabled)

src/modeling/onnx_export.py:41

↓ 1 callersFunctionbbox_xyxy2cs

Transform the bbox format from (x,y,w,h) into (center, scale) Args: bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted

src/dwpose/onnxpose.py:116

↓ 1 callersFunctionbetas_for_alpha_bar

Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [

src/scheduler/scheduler_lcm.py:54

↓ 1 callersMethodbind

(self, var_map, skip_check=SKIP_ENGINE_MODEL_CHECK)

src/modeling/engine_model.py:234

↓ 1 callersMethodclear_reference

(self)

src/morpher.py:444

↓ 1 callersFunctiondecode

Modulate simcc distribution with Gaussian. Args: simcc_x (np.ndarray[K, Wx]): model predicted simcc in x. simcc_y (np.ndarray[K,

src/dwpose/onnxpose.py:334

↓ 1 callersMethoddecode_frame

Decode a single latent timestep to RGB. Args: x: input NCHW latent (C=4) tensor with ~Gaussian values. mem: recurrent

src/taesdv/taesdv.py:68

↓ 1 callersMethoddecode_latents

(self, latents, vae_bsz, hiddens=None, push_ready=False)

src/pipeline/pipeline_pose2vid_lcm.py:104

↓ 1 callersMethoddecode_latents

(self, latents)

src/pipeline/pipeline_pose2vid.py:104

↓ 1 callersMethoddecode_slice

(self, vae, x, pmems)

src/modeling/framed_models.py:16

↓ 1 callersFunctiondemo_postprocess

(outputs, img_size, p6=False)

src/dwpose/onnxdet.py:58

↓ 1 callersFunctiondraw_bodypose

(canvas, candidate, subset)

src/dwpose/util.py:68

↓ 1 callersFunctiondraw_facepose

(canvas, all_lmks, translations=[], patterns=None)

src/dwpose/util.py:156

↓ 1 callersFunctiondraw_pose

(pose, H, W, translations, patterns, face_only)

src/dwpose/__init__.py:17

↓ 1 callersMethodfuse_reference

(self, ref_image)

src/morpher.py:463

↓ 1 callersMethodget_dynamic_axes

(self)

src/modeling/framed_models.py:100

↓ 1 callersFunctionget_engine

(engine_file_path)

src/modeling/engine_model.py:13

↓ 1 callersMethodget_input_names

(self)

src/modeling/framed_models.py:94

↓ 1 callersMethodget_output_names

(self)

src/modeling/framed_models.py:97

↓ 1 callersMethodget_processor

r""" Get the attention processor in use. Args: return_deprecated_lora (`bool`, *optional*, defaults to `False`):

src/models/attention_processor.py:406

↓ 1 callersMethodget_sample_input

(self, batchsize, height, width, dtype, device)

src/modeling/framed_models.py:51

↓ 1 callersMethodget_scalings_for_boundary_condition_discrete

(self, timestep)

src/scheduler/scheduler_lcm.py:380

↓ 1 callersFunctionget_simcc_maximum

Get maximum response location and value from simcc representations. Note: instance number: N num_keypoints: K heatmap hei

src/dwpose/onnxpose.py:289

↓ 1 callersFunctionget_warp_matrix

Calculate the affine transformation matrix that can warp the bbox area in the input image to the output size. Args: center (np.ndarra

src/dwpose/onnxpose.py:202

↓ 1 callersFunctioninference

Inference RTMPose model. Args: sess (ort.InferenceSession): ONNXRuntime session. img (np.ndarray): Input image in shape. Ret

src/dwpose/onnxpose.py:52

↓ 1 callersFunctioninference_detector

(session, oriImg)

src/dwpose/onnxdet.py:98

↓ 1 callersFunctioninference_pose

(session, out_bbox, oriImg)

src/dwpose/onnxpose.py:354

↓ 1 callersMethodload_reference_part_to_cpu_or_uninstall

(self)

src/morpher.py:418

↓ 1 callersMethodload_reference_part_to_device

(self, device=None)

src/morpher.py:432

↓ 1 callersFunctionmain

Run TAESDV roundtrip reconstruction on the given video paths.

src/taesdv/taesdv.py:167

↓ 1 callersFunctionmap_device

(device_or_str)

src/morpher.py:30

↓ 1 callersFunctionmulticlass_nms

Multiclass NMS implemented in Numpy. Class-aware version.

src/dwpose/onnxdet.py:35

↓ 1 callersFunctionnms

Single class NMS implemented in Numpy.

src/dwpose/onnxdet.py:6

↓ 1 callersFunctionnumpy_to_torch_dtype

(np_dtype)

src/modeling/engine_model.py:23

↓ 1 callersFunctionoptimize_onnx

( onnx_path: str, onnx_opt_path: str, )

src/modeling/onnx_export.py:66

↓ 1 callersMethodoutput

(self, dtype=torch.float16)

src/models/mutual_self_attention.py:398

↓ 1 callersFunctionparam2translations

(face_length_ratio=0.75, eye_height_ratio=3.0, eye_width_ratio=1.6, mouth_height_ratio=1.2, mouth_width_ratio=

pipeline_example.py:76

↓ 1 callersFunctionpostprocess

Postprocess for RTMPose model output. Args: outputs (np.ndarray): Output of RTMPose model. model_input_size (tuple): RTMPose mode

src/dwpose/onnxpose.py:81

↓ 1 callersMethodprepare_extra_step_kwargs

(self, generator, eta)

src/pipeline/pipeline_pose2vid_lcm.py:129

↓ 1 callersMethodprepare_extra_step_kwargs

(self, generator, eta)

src/pipeline/pipeline_pose2vid.py:119

↓ 1 callersMethodprepare_latents

( self, batch_size, num_channels_latents, width, height, video

src/pipeline/pipeline_pose2vid_lcm.py:150

↓ 1 callersMethodprepare_latents

( self, batch_size, num_channels_latents, width, height, video

src/pipeline/pipeline_pose2vid.py:140

↓ 1 callersFunctionpreprocess

(img, input_size, swap=(2, 0, 1))

src/dwpose/onnxdet.py:80

↓ 1 callersFunctionpreprocess

Do preprocessing for RTMPose model inference. Args: img (np.ndarray): Input image in shape. input_size (tuple): Input image size

src/dwpose/onnxpose.py:7

↓ 1 callersMethodprocess_input

(self, wait_after=0.0)

src/morpher.py:506

↓ 1 callersMethodregister_reference_hooks

( self, mode, do_classifier_free_guidance, attention_auto_machine_weight,

src/models/mutual_self_attention.py:55

↓ 1 callersFunctionrescale_zero_terminal_snr

Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1) Args: betas (`torch.FloatTenso

src/scheduler/scheduler_lcm.py:99

next →1–100 of 352, ranked by callers