|
|
import os, torch, json |
|
|
from .sd_video import ModelManager, SDVideoPipeline, ControlNetConfigUnit |
|
|
from ..processors.sequencial_processor import SequencialProcessor |
|
|
from ..data import VideoData, save_frames, save_video |
|
|
|
|
|
|
|
|
|
|
|
class SDVideoPipelineRunner: |
|
|
def __init__(self, in_streamlit=False): |
|
|
self.in_streamlit = in_streamlit |
|
|
|
|
|
|
|
|
def load_pipeline(self, model_list, textual_inversion_folder, device, lora_alphas, controlnet_units): |
|
|
|
|
|
model_manager = ModelManager(torch_dtype=torch.float16, device=device) |
|
|
model_manager.load_models(model_list) |
|
|
pipe = SDVideoPipeline.from_model_manager( |
|
|
model_manager, |
|
|
[ |
|
|
ControlNetConfigUnit( |
|
|
processor_id=unit["processor_id"], |
|
|
model_path=unit["model_path"], |
|
|
scale=unit["scale"] |
|
|
) for unit in controlnet_units |
|
|
] |
|
|
) |
|
|
textual_inversion_paths = [] |
|
|
for file_name in os.listdir(textual_inversion_folder): |
|
|
if file_name.endswith(".pt") or file_name.endswith(".bin") or file_name.endswith(".pth") or file_name.endswith(".safetensors"): |
|
|
textual_inversion_paths.append(os.path.join(textual_inversion_folder, file_name)) |
|
|
pipe.prompter.load_textual_inversions(textual_inversion_paths) |
|
|
return model_manager, pipe |
|
|
|
|
|
|
|
|
def load_smoother(self, model_manager, smoother_configs): |
|
|
smoother = SequencialProcessor.from_model_manager(model_manager, smoother_configs) |
|
|
return smoother |
|
|
|
|
|
|
|
|
def synthesize_video(self, model_manager, pipe, seed, smoother, **pipeline_inputs): |
|
|
torch.manual_seed(seed) |
|
|
if self.in_streamlit: |
|
|
import streamlit as st |
|
|
progress_bar_st = st.progress(0.0) |
|
|
output_video = pipe(**pipeline_inputs, smoother=smoother, progress_bar_st=progress_bar_st) |
|
|
progress_bar_st.progress(1.0) |
|
|
else: |
|
|
output_video = pipe(**pipeline_inputs, smoother=smoother) |
|
|
model_manager.to("cpu") |
|
|
return output_video |
|
|
|
|
|
|
|
|
def load_video(self, video_file, image_folder, height, width, start_frame_id, end_frame_id): |
|
|
video = VideoData(video_file=video_file, image_folder=image_folder, height=height, width=width) |
|
|
if start_frame_id is None: |
|
|
start_frame_id = 0 |
|
|
if end_frame_id is None: |
|
|
end_frame_id = len(video) |
|
|
frames = [video[i] for i in range(start_frame_id, end_frame_id)] |
|
|
return frames |
|
|
|
|
|
|
|
|
def add_data_to_pipeline_inputs(self, data, pipeline_inputs): |
|
|
pipeline_inputs["input_frames"] = self.load_video(**data["input_frames"]) |
|
|
pipeline_inputs["num_frames"] = len(pipeline_inputs["input_frames"]) |
|
|
pipeline_inputs["width"], pipeline_inputs["height"] = pipeline_inputs["input_frames"][0].size |
|
|
if len(data["controlnet_frames"]) > 0: |
|
|
pipeline_inputs["controlnet_frames"] = [self.load_video(**unit) for unit in data["controlnet_frames"]] |
|
|
return pipeline_inputs |
|
|
|
|
|
|
|
|
def save_output(self, video, output_folder, fps, config): |
|
|
os.makedirs(output_folder, exist_ok=True) |
|
|
save_frames(video, os.path.join(output_folder, "frames")) |
|
|
save_video(video, os.path.join(output_folder, "video.mp4"), fps=fps) |
|
|
config["pipeline"]["pipeline_inputs"]["input_frames"] = [] |
|
|
config["pipeline"]["pipeline_inputs"]["controlnet_frames"] = [] |
|
|
with open(os.path.join(output_folder, "config.json"), 'w') as file: |
|
|
json.dump(config, file, indent=4) |
|
|
|
|
|
|
|
|
def run(self, config): |
|
|
if self.in_streamlit: |
|
|
import streamlit as st |
|
|
if self.in_streamlit: st.markdown("Loading videos ...") |
|
|
config["pipeline"]["pipeline_inputs"] = self.add_data_to_pipeline_inputs(config["data"], config["pipeline"]["pipeline_inputs"]) |
|
|
if self.in_streamlit: st.markdown("Loading videos ... done!") |
|
|
if self.in_streamlit: st.markdown("Loading models ...") |
|
|
model_manager, pipe = self.load_pipeline(**config["models"]) |
|
|
if self.in_streamlit: st.markdown("Loading models ... done!") |
|
|
if "smoother_configs" in config: |
|
|
if self.in_streamlit: st.markdown("Loading smoother ...") |
|
|
smoother = self.load_smoother(model_manager, config["smoother_configs"]) |
|
|
if self.in_streamlit: st.markdown("Loading smoother ... done!") |
|
|
else: |
|
|
smoother = None |
|
|
if self.in_streamlit: st.markdown("Synthesizing videos ...") |
|
|
output_video = self.synthesize_video(model_manager, pipe, config["pipeline"]["seed"], smoother, **config["pipeline"]["pipeline_inputs"]) |
|
|
if self.in_streamlit: st.markdown("Synthesizing videos ... done!") |
|
|
if self.in_streamlit: st.markdown("Saving videos ...") |
|
|
self.save_output(output_video, config["data"]["output_folder"], config["data"]["fps"], config) |
|
|
if self.in_streamlit: st.markdown("Saving videos ... done!") |
|
|
if self.in_streamlit: st.markdown("Finished!") |
|
|
video_file = open(os.path.join(os.path.join(config["data"]["output_folder"], "video.mp4")), 'rb') |
|
|
if self.in_streamlit: st.video(video_file.read()) |
|
|
|
