DiffSynth-Studio/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py at main · modelscope/DiffSynth-Studio · GitHub

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
import torch
from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig
from diffsynth.utils.data.audio_video import write_video_audio
from diffsynth.utils.data import VideoData
from modelscope import dataset_snapshot_download

vram_config = {
    "offload_dtype": torch.bfloat16,
    "offload_device": "cpu",
    "onload_dtype": torch.bfloat16,
    "onload_device": "cpu",
    "preparing_dtype": torch.bfloat16,
    "preparing_device": "cuda",
    "computation_dtype": torch.bfloat16,
    "computation_device": "cuda",
}
dataset_snapshot_download(
    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
    local_dir="data/diffsynth_example_dataset",
    allow_file_pattern="minimax_h3/MiniMax-H3-FL2VA/*",
)
dataset_base_path = "data/diffsynth_example_dataset/minimax_h3/MiniMax-H3-FL2VA"
height, width, num_frames = 480, 832, 124
prompt = "A girl is very happy, she is speaking in english: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”"
frames = VideoData(f"{dataset_base_path}/video.mp4", height=height, width=width).raw_data()


def build_pipe(dit_path):
    return MiniMaxH3Pipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=[
            ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/text_encoder/model*.safetensors", **vram_config),
            ModelConfig(path=dit_path, **vram_config),
            ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/video_vae/source/model.safetensors", **vram_config),
            ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/audio_vae/model.safetensors", **vram_config),
        ],
        vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
    )


pipe = build_pipe("models/train/MiniMax-H3-T2VA-full/epoch-1.safetensors")
video, audio = pipe(
    prompt=prompt,
    height=height, width=width, num_frames=num_frames,
    num_inference_steps=50, seed=0,
)
write_video_audio(
    video=video, audio=audio, output_path="minimax_h3_t2va_full.mp4",
    fps=24, audio_sample_rate=pipe.audio_vae.sample_rate,
)
print("saved minimax_h3_t2va_full.mp4", "frames:", len(video), "audio:", tuple(audio.shape))
del pipe
torch.cuda.empty_cache()

pipe = build_pipe("models/train/MiniMax-H3-FL2VA-full/epoch-1.safetensors")
video, audio = pipe(
    prompt=prompt,
    height=height, width=width, num_frames=num_frames,
    num_inference_steps=50, seed=0,
    keyframes=[frames[0], frames[num_frames - 1]],
    keyframe_indices=[0, -1],
)
write_video_audio(
    video=video, audio=audio, output_path="minimax_h3_fl2va_full.mp4",
    fps=24, audio_sample_rate=pipe.audio_vae.sample_rate,
)
print("saved minimax_h3_fl2va_full.mp4", "frames:", len(video), "audio:", tuple(audio.shape))