This commit is contained in:
sayakpaul
2024-11-20 19:31:13 +05:30
parent 7e203db9aa
commit 9a15eaee41
2 changed files with 31 additions and 32 deletions
-1
View File
@@ -5,7 +5,6 @@ import numpy as np
import torch import torch
import torchvision.transforms as TT import torchvision.transforms as TT
from accelerate.logging import get_logger from accelerate.logging import get_logger
from torchvision import transforms
from torchvision.transforms import InterpolationMode from torchvision.transforms import InterpolationMode
from torchvision.transforms.functional import resize from torchvision.transforms.functional import resize
+5 -5
View File
@@ -76,7 +76,7 @@ def save_model_card(
fps=8, fps=8,
): ):
widget_dict = [] widget_dict = []
if videos is not None: if videos is not None and len(videos) > 0:
for i, video in enumerate(videos): for i, video in enumerate(videos):
export_to_video(video, os.path.join(repo_folder, f"final_video_{i}.mp4", fps=fps)) export_to_video(video, os.path.join(repo_folder, f"final_video_{i}.mp4", fps=fps))
widget_dict.append( widget_dict.append(
@@ -876,11 +876,13 @@ def main(args):
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.synchronize(accelerator.device) torch.cuda.synchronize(accelerator.device)
# Final test inference
validation_outputs = []
if args.validation_prompt and args.num_validation_videos > 0:
accelerator.print("===== Memory before testing =====") accelerator.print("===== Memory before testing =====")
print_memory(accelerator.device) print_memory(accelerator.device)
reset_memory(accelerator.device) reset_memory(accelerator.device)
# Final test inference
pipe = MochiPipeline.from_pretrained( pipe = MochiPipeline.from_pretrained(
args.pretrained_model_name_or_path, args.pretrained_model_name_or_path,
revision=args.revision, revision=args.revision,
@@ -902,8 +904,6 @@ def main(args):
pipe.set_adapters(["mochi-lora"], [lora_scaling]) pipe.set_adapters(["mochi-lora"], [lora_scaling])
# Run inference # Run inference
validation_outputs = []
if args.validation_prompt and args.num_validation_videos > 0:
validation_prompts = args.validation_prompt.split(args.validation_prompt_separator) validation_prompts = args.validation_prompt.split(args.validation_prompt_separator)
for validation_prompt in validation_prompts: for validation_prompt in validation_prompts:
pipeline_args = { pipeline_args = {
@@ -942,7 +942,7 @@ def main(args):
repo_id=repo_id, repo_id=repo_id,
folder_path=args.output_dir, folder_path=args.output_dir,
commit_message="End of training", commit_message="End of training",
ignore_patterns=["step_*", "epoch_*", "*.bin"], ignore_patterns=["step_*", "epoch_*", "*.bin", "*.pt"],
) )
accelerator.end_training() accelerator.end_training()