File size: 2,824 Bytes
e14f899 |
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 |
import json
import numpy as np
import imageio
import os
import torch
import torch.distributed as dist
def export_to_video(video_frames, output_video_path, fps = 12):
# Ensure all frames are NumPy arrays and determine video dimensions from the first frame
assert all(isinstance(frame, np.ndarray) for frame in video_frames), "All video frames must be NumPy arrays."
# Ensure output_video_path is ending with .mp4
if not output_video_path.endswith('.mp4'):
output_video_path += '.mp4'
# Create a video file at the specified path and write frames to it
with imageio.get_writer(output_video_path, fps=fps, format='mp4') as writer:
for frame in video_frames:
writer.append_data(
(frame * 255).astype(np.uint8)
)
def save_generation(video_frames, configs, base_path, file_name=None):
if not os.path.exists(base_path):
os.makedirs(base_path)
p_config = configs["pipe_configs"]
frames, steps, fps = p_config["num_frames"], p_config["steps"], p_config["fps"]
if not file_name:
index = [int(each.split('_')[0]) for each in os.listdir(base_path)]
max_idex = max(index) if index else 0
idx_str = str(max_idex + 1).zfill(6)
key_info = '_'.join([str(frames), str(steps), str(fps)])
file_name = f'{idx_str}_{key_info}'
with open(f'{base_path}/{file_name}.json', 'w') as f:
json.dump(configs, f, indent=4)
export_to_video(video_frames, f'{base_path}/{file_name}.mp4', fps=p_config["export_fps"])
return file_name
class GlobalState:
def __init__(self, state={}) -> None:
self.init_state(state)
def init_state(self, state={}):
self.state = state
def set(self, key, value):
self.state[key] = value
def get(self, key, default=None):
return self.state.get(key, default)
class DistController(object):
def __init__(self, rank, world_size, config = None) -> None:
super().__init__()
self.rank = rank
self.world_size = world_size
self.config = config
self.is_master = (rank == 0)
print("DistController is master: ", self.is_master)
#self.init_dist()
self.init_group()
#self.device = torch.device(f"cuda:{config['devices'][dist.get_rank()]}")
self.device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(self.device)
def init_dist(self):
print(f"Rank {self.rank} is running.")
os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = str(self.config.get("master_port") or "29500")
dist.init_process_group("nccl", rank=self.rank, world_size=self.world_size)
def init_group(self):
self.adj_groups = [dist.new_group([i, i+1]) for i in range(self.world_size-1)]
|