diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py index 83d3497d4..2d7381427 100644 --- a/diffsynth/diffusion/flow_match.py +++ b/diffsynth/diffusion/flow_match.py @@ -208,6 +208,7 @@ def set_timesteps_z_image(num_inference_steps=100, denoising_strength=1.0, shift for timestep in target_timesteps: timestep_id = torch.argmin((timesteps - timestep).abs()) timesteps[timestep_id] = timestep + sigmas[timestep_id] = timestep / num_train_timesteps return sigmas, timesteps @staticmethod diff --git a/tests/test_flow_match.py b/tests/test_flow_match.py new file mode 100644 index 000000000..54f18ef25 --- /dev/null +++ b/tests/test_flow_match.py @@ -0,0 +1,26 @@ +import importlib.util +from pathlib import Path + +import torch + +module_path = Path(__file__).parents[1] / "diffsynth" / "diffusion" / "flow_match.py" +module_spec = importlib.util.spec_from_file_location("flow_match", module_path) +flow_match = importlib.util.module_from_spec(module_spec) +module_spec.loader.exec_module(flow_match) + + +def test_z_image_target_timesteps_keep_matching_sigmas(): + student = flow_match.FlowMatchScheduler("Z-Image") + student.set_timesteps(8) + + teacher = flow_match.FlowMatchScheduler("Z-Image") + teacher.set_timesteps(50, target_timesteps=student.timesteps) + + for target_timestep in student.timesteps: + timestep_id = torch.argmin((teacher.timesteps - target_timestep).abs()) + torch.testing.assert_close(teacher.timesteps[timestep_id], target_timestep) + + torch.testing.assert_close( + teacher.timesteps, + teacher.sigmas * teacher.num_train_timesteps, + )