(self, hidden_states, device, tile_size=(34, 34), tile_stride=(16, 16))
| 1050 | return mask |
| 1051 | |
| 1052 | def tiled_decode(self, hidden_states, device, tile_size=(34, 34), tile_stride=(16, 16)): |
| 1053 | B, T, C, H, W = hidden_states.shape |
| 1054 | size_h, size_w = tile_size |
| 1055 | stride_h, stride_w = tile_stride |
| 1056 | |
| 1057 | # Split tasks |
| 1058 | tasks = [] |
| 1059 | for t in range(0, T, 3): |
| 1060 | for h in range(0, H, stride_h): |
| 1061 | if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue |
| 1062 | for w in range(0, W, stride_w): |
| 1063 | if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue |
| 1064 | t_, h_, w_ = t + 3, h + size_h, w + size_w |
| 1065 | tasks.append((t, t_, h, h_, w, w_)) |
| 1066 | |
| 1067 | # Run |
| 1068 | data_device = "cpu" |
| 1069 | computation_device = device |
| 1070 | |
| 1071 | weight = torch.zeros((1, 1, T//3*17, H * 16, W * 16), dtype=hidden_states.dtype, device=data_device) |
| 1072 | values = torch.zeros((B, 3, T//3*17, H * 16, W * 16), dtype=hidden_states.dtype, device=data_device) |
| 1073 | |
| 1074 | for t, t_, h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): |
| 1075 | hidden_states_batch = hidden_states[:, t:t_, :, h:h_, w:w_].to(computation_device) |
| 1076 | hidden_states_batch = self.decode_naive(hidden_states_batch, True).to(data_device) |
| 1077 | |
| 1078 | mask = self.build_mask( |
| 1079 | hidden_states_batch, |
| 1080 | is_bound=(h==0, h_>=H, w==0, w_>=W), |
| 1081 | border_width=((size_h - stride_h) * 16, (size_w - stride_w) * 16) |
| 1082 | ).to(dtype=hidden_states.dtype, device=data_device) |
| 1083 | |
| 1084 | target_t = t // 3 * 17 |
| 1085 | target_h = h * 16 |
| 1086 | target_w = w * 16 |
| 1087 | values[ |
| 1088 | :, |
| 1089 | :, |
| 1090 | target_t: target_t + hidden_states_batch.shape[2], |
| 1091 | target_h: target_h + hidden_states_batch.shape[3], |
| 1092 | target_w: target_w + hidden_states_batch.shape[4], |
| 1093 | ] += hidden_states_batch * mask |
| 1094 | weight[ |
| 1095 | :, |
| 1096 | :, |
| 1097 | target_t: target_t + hidden_states_batch.shape[2], |
| 1098 | target_h: target_h + hidden_states_batch.shape[3], |
| 1099 | target_w: target_w + hidden_states_batch.shape[4], |
| 1100 | ] += mask |
| 1101 | return values / weight |
| 1102 | |
| 1103 | def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(16, 16), smooth_scale=0.6): |
| 1104 | hidden_states = hidden_states.to("cpu") |
no test coverage detected