fix postprocess in distillation val

This commit is contained in:
H1yori233
2026-03-18 17:39:42 -07:00
parent 9380392110
commit e16cce72ad
2 changed files with 16 additions and 15 deletions
@@ -1318,6 +1318,7 @@ class DistillationPipeline(TrainingPipeline):
x = torchvision.utils.make_grid(x, nrow=6)
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
frames.append((x * 255).numpy().astype(np.uint8))
frames = self._post_process_validation_frames(frames, batch)
videos.append(frames)
audios.append(output_batch.extra.get("audio"))
audio_sample_rates.append(
@@ -142,18 +142,18 @@ class MatrixGameSelfForcingDistillationPipeline(SelfForcingDistillationPipeline
getattr(training_args, "warp_denoising_step", False)
or getattr(pipeline_config, "warp_denoising_step", False)
)
training_args.warp_denoising_step = warp_denoising_step
pipeline_config.warp_denoising_step = warp_denoising_step
self.use_diagonal_denoising = bool(
getattr(pipeline_config, "use_diagonal_denoising", False)
)
pipeline_config.use_diagonal_denoising = self.use_diagonal_denoising
raw_warmup_mid_steps = getattr(
pipeline_config, "diagonal_warmup_mid_steps", None
)
pipeline_config.diagonal_warmup_mid_steps = raw_warmup_mid_steps
training_args.warp_denoising_step = warp_denoising_step
pipeline_config.warp_denoising_step = warp_denoising_step
self.use_diagonal_denoising = bool(
getattr(pipeline_config, "use_diagonal_denoising", False)
)
pipeline_config.use_diagonal_denoising = self.use_diagonal_denoising
raw_warmup_mid_steps = getattr(
pipeline_config, "diagonal_warmup_mid_steps", None
)
pipeline_config.diagonal_warmup_mid_steps = raw_warmup_mid_steps
self.denoising_step_list = resolve_denoising_steps(
pipeline_config.dmd_denoising_steps,
@@ -1356,13 +1356,13 @@ class MatrixGameSelfForcingDistillationPipeline(SelfForcingDistillationPipeline
key_names[i]: bool(keyboard_cond[frame_idx, i])
for i in range(min(len(key_names), keyboard_cond.shape[1]))
}
draw_keys_on_frame(frame, keys, mode='universal')
draw_keys_on_frame(frame, keys)
# Draw mouse overlay
if mouse_cond is not None and frame_idx < len(mouse_cond):
pitch = float(mouse_cond[frame_idx, 0])
yaw = float(mouse_cond[frame_idx, 1])
draw_mouse_on_frame(frame, pitch, yaw)
draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
processed_frames.append(frame)
@@ -1388,4 +1388,4 @@ if __name__ == "__main__":
parser = TrainingArgs.add_cli_args(parser)
parser = FastVideoArgs.add_cli_args(parser)
args = parser.parse_args()
main(args)
main(args)