storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-22 08:05:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

updates.

Browse Source
This commit is contained in:
sayakpaul
2024-11-19 10:21:28 +05:30
parent ac83c78a5a
commit 1409d478d5
5 changed files with 22 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/mochi-1/prepare_dataset.py
+5 -4
View File
@@ -321,8 +321,11 @@ def serialize_artifacts(
prompt_embeds: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None
) -> None:
num_frames, height, width = videos.size(1), videos.size(3), videos.size(4)
metadata = [{"num_frames": num_frames, "height": height, "width": width}]
metadata = []
for i in range(videos.size(0)):
video = videos[i:i+1]
metadata_dict = {"num_frames": video.size(1), "height": video.size(3), "width": video.size(4)}
metadata.append(metadata_dict)
data_folder_mapper_list = [
(images, images_dir, lambda img, path: save_image(img[0], path), "png"),
@@ -542,7 +545,6 @@ def main():
video_latents = vae._encode(videos)
video_latents = video_latents.to(memory_format=torch.contiguous_format, dtype=weight_dtype)
print(f"{video_latents.shape=}")
# Encode prompts
prompt_embeds, prompt_attention_mask = compute_prompt_embeddings(
@@ -554,7 +556,6 @@ def main():
weight_dtype,
requires_grad=False,
)
print(f"{prompt_attention_mask.shape=}")
if images is not None:
images = (images.permute(0, 2, 1, 3, 4) + 1) / 2