storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-11-02 18:06:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Windows support for T2V scripts (#48)

Browse Source
This commit is contained in:
Aryan
2024-10-21 02:40:05 +05:30
committed by GitHub
parent 6c00cf094b
commit db9b295408
3 changed files with 46 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/cogvideox_image_to_video_lora.py
+4 -4
View File
@@ -202,11 +202,11 @@ def log_validation(
class CollateFunction:
def __init__(self, weight_dtype, load_tensors):
def __init__(self, weight_dtype: torch.dtype, load_tensors: bool) -> None:
self.weight_dtype = weight_dtype
self.load_tensors = load_tensors
def __call__(self, data):
def __call__(self, data: Dict[str, Any]) -> Dict[str, torch.Tensor]:
prompts = [x["prompt"] for x in data[0]]
if self.load_tensors:
@@ -519,13 +519,13 @@ def main(args):
video_reshape_mode=args.video_reshape_mode, **dataset_init_kwargs
)
collate_fn_instance = CollateFunction(weight_dtype, args.load_tensors)
collate_fn = CollateFunction(weight_dtype, args.load_tensors)
train_dataloader = DataLoader(
train_dataset,
batch_size=1,
sampler=BucketSampler(train_dataset, batch_size=args.train_batch_size, shuffle=True),
collate_fn=collate_fn_instance,
collate_fn=collate_fn,
num_workers=args.dataloader_num_workers,
pin_memory=args.pin_memory,
)