File size: 4,251 Bytes
74da989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | from typing import Any
from trainer.pipelines.pipeline_batch_info import PreprocessBatch
def basic_t2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
"""Create a record for the Parquet dataset from PreprocessBatch."""
# For batch processing, we need to handle the case where some fields might be single values
# or lists depending on the batch size
assert isinstance(batch.prompt, list)
assert isinstance(batch.width, list)
assert isinstance(batch.height, list)
assert isinstance(batch.fps, list)
assert isinstance(batch.num_frames, list)
records = []
for idx, video_name in enumerate(batch.video_file_name):
width = batch.width[idx] if batch.width is not None else 0
height = batch.height[idx] if batch.height is not None else 0
# Get FPS - single value in PreprocessBatch
fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0
# For duration, we need to calculate it or use a default since it's not in PreprocessBatch
# duration = num_frames / fps if available
duration_val = 0.0
if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0:
duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx])
record = {
"id":
video_name,
"vae_latent_bytes":
batch.latents[idx].tobytes(),
"vae_latent_shape":
list(batch.latents[idx].shape),
"vae_latent_dtype":
str(batch.latents[idx].dtype),
"text_embedding_bytes":
batch.prompt_embeds[idx].tobytes(),
"text_embedding_shape":
list(batch.prompt_embeds[idx].shape),
"text_embedding_dtype":
str(batch.prompt_embeds[idx].dtype),
"file_name":
video_name,
"caption":
batch.prompt[idx],
"media_type":
"video",
"width":
int(width),
"height":
int(height),
"num_frames":
batch.latents[idx].shape[1]
if len(batch.latents[idx].shape) > 1 else 0,
"duration_sec":
duration_val,
"fps":
fps_val,
}
records.append(record)
return records
def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
"""Create a record for the Parquet dataset with CLIP features."""
records = basic_t2v_record_creator(batch)
assert len(
batch.image_embeds) == 1, "image embedding should be a single tensor"
image_embeds = batch.image_embeds[0]
image_latent = batch.image_latent
pil_image = batch.pil_image
for idx, record in enumerate(records):
if image_embeds is not None:
record.update({
"clip_feature_bytes": image_embeds[idx].tobytes(),
"clip_feature_shape": list(image_embeds[idx].shape),
"clip_feature_dtype": str(image_embeds[idx].dtype),
})
else:
record.update({
"clip_feature_bytes": b"",
"clip_feature_shape": [],
"clip_feature_dtype": "",
})
if image_latent is not None:
record.update({
"first_frame_latent_bytes":
image_latent[idx].tobytes(),
"first_frame_latent_shape":
list(image_latent[idx].shape),
"first_frame_latent_dtype":
str(image_latent[idx].dtype),
})
else:
record.update({
"first_frame_latent_bytes": b"",
"first_frame_latent_shape": [],
"first_frame_latent_dtype": "",
})
if pil_image is not None:
record.update({
"pil_image_bytes": pil_image[idx].tobytes(),
"pil_image_shape": list(pil_image[idx].shape),
"pil_image_dtype": str(pil_image[idx].dtype),
})
else:
record.update({
"pil_image_bytes": b"",
"pil_image_shape": [],
"pil_image_dtype": "",
})
return records
|