| """Sliding-window decode helpers used by the Gradio demo.""" |
|
|
| import torch |
| from torch_geometric.data import Batch |
|
|
|
|
| def decode_to_hatD_with_sliding_window( |
| model, |
| latent_codes, |
| src_graphs_list, |
| total_frames, |
| window_size, |
| overlap, |
| device, |
| model_type, |
| ): |
| """Decode a long latent/code sequence in overlapping windows. |
| |
| Overlapping regions keep the earlier window prediction and drop the later |
| window overlap. This matches the preprocessing convention used by the demo |
| checkpoints. |
| """ |
| if model_type not in {"vae", "rvq"}: |
| raise ValueError(f"model_type must be 'vae' or 'rvq', got: {model_type}") |
|
|
| stride = window_size - overlap |
| if stride <= 0: |
| raise ValueError("window_size must be greater than overlap") |
|
|
| if len(src_graphs_list) != total_frames: |
| raise ValueError( |
| f"src_graphs_list length ({len(src_graphs_list)}) != total_frames ({total_frames})" |
| ) |
| if latent_codes.shape[0] != total_frames: |
| raise ValueError( |
| f"latent_codes length ({latent_codes.shape[0]}) != total_frames ({total_frames})" |
| ) |
|
|
| if total_frames <= window_size: |
| num_windows = 1 |
| else: |
| num_windows = (total_frames - window_size + stride - 1) // stride + 1 |
|
|
| num_nodes_per_frame = src_graphs_list[0].skel_x.shape[0] |
| hatD_parts = [] |
|
|
| for window_idx in range(num_windows): |
| start_frame = window_idx * stride |
| if start_frame >= total_frames: |
| break |
|
|
| end_frame = min(start_frame + window_size, total_frames) |
| window_frames = end_frame - start_frame |
| latent_window = latent_codes[start_frame:end_frame].to(device) |
|
|
| src_window_graphs = src_graphs_list[start_frame:end_frame] |
| src_window = Batch.from_data_list(src_window_graphs).to(device) |
|
|
| with torch.no_grad(): |
| if model_type == "vae": |
| hatD_win = model.decode(latent_window, src_window, window_frames) |
| else: |
| hatD_win, _ = model.decode_from_codes(latent_window, src_window, window_frames) |
|
|
| hatD_dim = hatD_win.shape[1] |
| hatD_win_reshaped = hatD_win.view(window_frames, num_nodes_per_frame, hatD_dim) |
|
|
| if window_idx == 0: |
| keep_start_idx = 0 |
| else: |
| keep_start_idx = overlap |
| hatD_parts.append(hatD_win_reshaped[keep_start_idx:window_frames].cpu()) |
|
|
| del src_window, hatD_win, hatD_win_reshaped |
|
|
| hatD_full_3d = torch.cat(hatD_parts, dim=0) |
| actual_frames = hatD_full_3d.shape[0] |
| hatD_full = hatD_full_3d.view(actual_frames * num_nodes_per_frame, -1).to(device) |
| src_batch_full = Batch.from_data_list(src_graphs_list).to(device) |
|
|
| return hatD_full, src_batch_full, actual_frames, num_nodes_per_frame |
|
|
