mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-09 00:09:45 +00:00
sayakpaul
112 lines
4.2 KiB
Python
112 lines
4.2 KiB
Python
"""
|
|
Adapted from:
|
|
https://github.com/genmoai/mochi/blob/main/demos/fine_tuner/encode_videos.py
|
|
https://github.com/genmoai/mochi/blob/main/demos/fine_tuner/embed_captions.py
|
|
"""
|
|
|
|
import click
|
|
import torch
|
|
import torchvision
|
|
from pathlib import Path
|
|
from diffusers import AutoencoderKLMochi, MochiPipeline
|
|
from transformers import T5EncoderModel, T5Tokenizer
|
|
from tqdm.auto import tqdm
|
|
|
|
|
|
def encode_videos(model: torch.nn.Module, vid_path: Path, shape: str):
|
|
T, H, W = [int(s) for s in shape.split("x")]
|
|
assert (T - 1) % 6 == 0, "Expected T to be 1 mod 6"
|
|
video, _, metadata = torchvision.io.read_video(str(vid_path), output_format="THWC", pts_unit="secs")
|
|
fps = metadata["video_fps"]
|
|
video = video.permute(3, 0, 1, 2)
|
|
og_shape = video.shape
|
|
assert video.shape[2] == H, f"Expected {vid_path} to have height {H}, got {video.shape}"
|
|
assert video.shape[3] == W, f"Expected {vid_path} to have width {W}, got {video.shape}"
|
|
assert video.shape[1] >= T, f"Expected {vid_path} to have at least {T} frames, got {video.shape}"
|
|
if video.shape[1] > T:
|
|
video = video[:, :T]
|
|
print(f"Trimmed video from {og_shape[1]} to first {T} frames")
|
|
video = video.unsqueeze(0)
|
|
video = video.float() / 127.5 - 1.0
|
|
video = video.to(model.device)
|
|
|
|
assert video.ndim == 5
|
|
|
|
with torch.inference_mode():
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
ldist = model._encode(video)
|
|
|
|
torch.save(dict(ldist=ldist), vid_path.with_suffix(".latent.pt"))
|
|
|
|
|
|
@click.command()
|
|
@click.argument("output_dir", type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path))
|
|
@click.option(
|
|
"--model_id",
|
|
type=str,
|
|
help="Repo id. Should be genmo/mochi-1-preview",
|
|
default="genmo/mochi-1-preview",
|
|
)
|
|
@click.option("--shape", default="163x480x848", help="Shape of the video to encode")
|
|
@click.option("--overwrite", "-ow", is_flag=True, help="Overwrite existing latents and caption embeddings.")
|
|
def batch_process(output_dir: Path, model_id: Path, shape: str, overwrite: bool) -> None:
|
|
"""Process all videos and captions in a directory using a single GPU."""
|
|
# comment out when running on unsupported hardware
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
|
|
# Get all video paths
|
|
video_paths = list(output_dir.glob("**/*.mp4"))
|
|
if not video_paths:
|
|
print(f"No MP4 files found in {output_dir}")
|
|
return
|
|
|
|
text_paths = list(output_dir.glob("**/*.txt"))
|
|
if not text_paths:
|
|
print(f"No text files found in {output_dir}")
|
|
return
|
|
|
|
# load the models
|
|
vae = AutoencoderKLMochi.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32).to("cuda")
|
|
text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder")
|
|
tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
|
|
pipeline = MochiPipeline.from_pretrained(
|
|
model_id, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None, vae=None
|
|
).to("cuda")
|
|
|
|
for idx, video_path in tqdm(enumerate(sorted(video_paths))):
|
|
print(f"Processing {video_path}")
|
|
try:
|
|
if video_path.with_suffix(".latent.pt").exists() and not overwrite:
|
|
print(f"Skipping {video_path}")
|
|
continue
|
|
|
|
# encode videos.
|
|
encode_videos(vae, vid_path=video_path, shape=shape)
|
|
|
|
# embed captions.
|
|
prompt_path = Path("/".join(str(video_path).split(".")[:-1]) + ".txt")
|
|
embed_path = prompt_path.with_suffix(".embed.pt")
|
|
|
|
if embed_path.exists() and not overwrite:
|
|
print(f"Skipping {prompt_path} - embeddings already exist")
|
|
continue
|
|
|
|
with open(prompt_path) as f:
|
|
text = f.read().strip()
|
|
with torch.inference_mode():
|
|
conditioning = pipeline.encode_prompt(prompt=[text])
|
|
|
|
conditioning = {"prompt_embeds": conditioning[0], "prompt_attention_mask": conditioning[1]}
|
|
torch.save(conditioning, embed_path)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
print(f"Error processing {video_path}: {str(e)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
batch_process()
|