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

Update train_wan_t2v.py

Browse Source
This commit is contained in:
Zhongjie Duan
2025-04-17 15:37:30 +08:00
committed by GitHub
parent b36cad6929
commit bf81de0c88
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/wanvideo/train_wan_t2v.py
View File

@@ -56,13 +56,16 @@ class TextVideoDataset(torch.utils.data.Dataset):
frame = Image.fromarray(frame) frame = Image.fromarray(frame)
frame = self.crop_and_resize(frame) frame = self.crop_and_resize(frame)
if first_frame is None: if first_frame is None:
first_frame = np.array(frame) first_frame = frame
frame = frame_process(frame) frame = frame_process(frame)
frames.append(frame) frames.append(frame)
reader.close() reader.close()
frames = torch.stack(frames, dim=0) frames = torch.stack(frames, dim=0)
frames = rearrange(frames, "T C H W -> C T H W") frames = rearrange(frames, "T C H W -> C T H W")
first_frame = v2.functional.center_crop(first_frame, output_size=(self.height, self.width))
first_frame = np.array(first_frame)
if self.is_i2v: if self.is_i2v:
return frames, first_frame return frames, first_frame