theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-23 09:28:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

support wan tensor parallel (preview)

Browse Source
This commit is contained in:
Artiprocher
2025-03-17 19:39:45 +08:00
parent 39890f023f
commit 04d03500ff
5 changed files with 147 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
diffsynth/schedulers/flow_match.py
View File

@@ -37,7 +37,7 @@ class FlowMatchScheduler():
self.linear_timesteps_weights = bsmntw_weighing
def step(self, model_output, timestep, sample, to_final=False):
def step(self, model_output, timestep, sample, to_final=False, **kwargs):
if isinstance(timestep, torch.Tensor):
timestep = timestep.cpu()
timestep_id = torch.argmin((self.timesteps - timestep).abs())