Spaces:
Runtime error
Runtime error
File size: 2,769 Bytes
0328207 | 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 | import argparse
from pathlib import Path
import torch
from hydra import compose, initialize
from hydra.core.global_hydra import GlobalHydra
from torchvision.io import write_video
from optimize_utils import MultiTrajectory
from stream_drag_inference_wrapper import StreamDragInferenceWrapper
from utils.misc import set_seed
from video_operations import run_optimization, save_videos
def build_stream_drag_inference(
config_dir: str,
config_name: str,
checkpoint_path: str,
total_generate_block_number: int,
use_ema: bool,
seed: int,
) -> StreamDragInferenceWrapper:
if GlobalHydra.instance().is_initialized():
GlobalHydra.instance().clear()
with initialize(version_base=None, config_path=config_dir):
stream_config = compose(config_name=config_name)
return StreamDragInferenceWrapper(
stream_model_config=stream_config,
checkpoint_path=checkpoint_path,
total_generate_block_number=total_generate_block_number,
use_ema=use_ema,
seed=seed,
)
def main() -> None:
prompt_index = 4
trajectory_dir = "./saved_labels/self_forcing_dmd_vsink_stream_drag-seed42/0004-A_close-up_3D_animated_scene_of_a_short,_fluffy_mo"
start_block_index = 3
trajectory_prefix = "block_3_Animation"
config_dir = "configs"
config_name = "self_forcing_dmd_vsink_stream_drag"
checkpoint_path = "./checkpoints/self_forcing_dmd.pt"
total_generate_block_number = 36
seed = 42
fps = 8
output_dir = "outputs-editing/self_forcing_dmd_vsink_stream_drag-seed42"
use_ema = True
torch.set_grad_enabled(False)
trajectory = MultiTrajectory.load(
save_dir=trajectory_dir,
prefix=trajectory_prefix,
)
model = build_stream_drag_inference(
config_dir=config_dir,
config_name=config_name,
checkpoint_path=checkpoint_path,
total_generate_block_number=total_generate_block_number,
use_ema=use_ema,
seed=seed,
)
set_seed(seed)
model.reset()
model.inference(
start_block_index=0,
end_block_index=start_block_index,
prompt=trajectory.prompt,
)
all_video, current_video, end_block_index = run_optimization(
model=model,
trajectory=trajectory,
start_block_index=start_block_index,
)
saved_path = save_videos(
all_video=all_video,
current_video=current_video,
output_dir=Path(output_dir),
prompt_index=prompt_index,
prompt=trajectory.prompt,
start_block_index=start_block_index,
end_block_index=end_block_index,
mode=trajectory.drag_or_animation_select,
fps=fps,
)
print(str(saved_path))
if __name__ == "__main__":
main()
|