|
"""Command line interface for generating videos from the model.""" |
|
from __future__ import annotations |
|
|
|
import argparse |
|
import queue |
|
import threading |
|
|
|
import rerun as rr |
|
|
|
import vista |
|
|
|
|
|
def generate_local( |
|
first_frame_file_name: str, |
|
height=576, |
|
width=1024, |
|
n_rounds=4, |
|
n_frames=25, |
|
n_steps=10, |
|
cfg_scale=2.5, |
|
cond_aug=0.0, |
|
): |
|
|
|
log_queue = queue.SimpleQueue() |
|
|
|
handle = threading.Thread( |
|
target=vista.run_sampling, |
|
args=[ |
|
log_queue, |
|
first_frame_file_name, |
|
height, |
|
width, |
|
n_rounds, |
|
n_frames, |
|
n_steps, |
|
cfg_scale, |
|
cond_aug, |
|
], |
|
) |
|
handle.start() |
|
while True: |
|
msg = log_queue.get() |
|
if msg == "done": |
|
break |
|
else: |
|
entity_path, entity, times = msg |
|
rr.reset_time() |
|
for timeline, time in times: |
|
if isinstance(time, int): |
|
rr.set_time_sequence(timeline, time) |
|
else: |
|
rr.set_time_seconds(timeline, time) |
|
rr.log(entity_path, entity) |
|
handle.join() |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser( |
|
description="Generate video conditioned on a single image using the Vista model." |
|
) |
|
parser.add_argument( |
|
"--img-path", |
|
type=str, |
|
help="Path to image used as input for Canny edge detector.", |
|
default="./example_images/nus-0.jpg", |
|
) |
|
parser.add_argument( |
|
"--num-steps", |
|
type=int, |
|
help="Number of diffusion steps per image. Recommended range: 10-50. Higher values result in more detailed images and less blurry results.", |
|
default=20, |
|
) |
|
parser.add_argument( |
|
"--num-segments", |
|
type=int, |
|
help="Number of segments to generate. Each segment consists of 25 frames.", |
|
default=3, |
|
) |
|
rr.script_add_args(parser) |
|
args = parser.parse_args() |
|
rr.script_setup( |
|
args, |
|
"rerun_example_vista", |
|
default_blueprint=vista.generate_blueprint(args.num_segments), |
|
) |
|
|
|
generate_local(args.img_path, n_steps=args.num_steps, n_rounds=args.num_segments) |
|
|
