fix postprocess in distillation val
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user