99 lines
4.9 KiB
Python
99 lines
4.9 KiB
Python
from base64 import b64encode
|
|
from tqdm import tqdm
|
|
from omegaconf import OmegaConf
|
|
import time, os
|
|
from t2v_helpers.general_utils import get_t2v_version
|
|
from t2v_helpers.args import get_outdir, process_args
|
|
import modules.paths as ph
|
|
import t2v_helpers.args as t2v_helpers_args
|
|
from modules.shared import state
|
|
|
|
# VideoCrafter support is heavy WIP and sketchy, needs help and more devs!
|
|
def process_videocrafter(args_dict):
|
|
args, video_args = process_args(args_dict)
|
|
print(f"\033[4;33m text2video extension for auto1111 webui\033[0m")
|
|
print(f"Git commit: {get_t2v_version()}")
|
|
init_timestring = time.strftime('%Y%m%d%H%M%S')
|
|
outdir_current = os.path.join(get_outdir(), f"{init_timestring}")
|
|
|
|
os.makedirs(outdir_current, exist_ok=True)
|
|
|
|
# load & merge config
|
|
|
|
config_path = os.path.join(ph.models_path, "models/VideoCrafter/model_config.yaml")
|
|
if not os.path.exists(config_path):
|
|
config_path = os.path.join(os.getcwd(), "extensions/sd-webui-modelscope-text2video/scripts/videocrafter/base_t2v/model_config.yaml")
|
|
if not os.path.exists(config_path):
|
|
config_path = os.path.join(os.getcwd(), "extensions/sd-webui-text2video/scripts/videocrafter/base_t2v/model_config.yaml")
|
|
if not os.path.exists(config_path):
|
|
raise FileNotFoundError(f'Could not find config file at {os.path.join(ph.models_path, "models/VideoCrafter/model_config.yaml")}, nor at {os.path.join(os.getcwd(), "extensions/sd-webui-modelscope-text2video/scripts/videocrafter/base_t2v/model_config.yaml")}, nor at {os.path.join(os.getcwd(), "extensions/sd-webui-text2video/scripts/videocrafter/base_t2v/model_config.yaml")}')
|
|
|
|
config = OmegaConf.load(config_path)
|
|
print("VideoCrafter config: \n", config)
|
|
|
|
from videocrafter.lvdm.samplers.ddim import DDIMSampler
|
|
from videocrafter.sample_utils import load_model, get_conditions, make_model_input_shape, torch_to_np
|
|
from videocrafter.sample_text2video import sample_text2video
|
|
from videocrafter.lvdm.utils.saving_utils import npz_to_video_grid
|
|
from t2v_helpers.video_audio_utils import add_soundtrack
|
|
|
|
# get model & sampler
|
|
model, _, _ = load_model(config, ph.models_path+'/VideoCrafter/model.ckpt', #TODO: support safetensors and stuff
|
|
inject_lora=False, # TODO
|
|
lora_scale=1, # TODO
|
|
lora_path=ph.models_path+'/VideoCrafter/LoRA/LoRA.ckpt', #TODO: support LoRA and stuff
|
|
)
|
|
ddim_sampler = DDIMSampler(model)# if opt.sample_type == "ddim" else None
|
|
|
|
# if opt.inject_lora:
|
|
# assert(opt.lora_trigger_word != '')
|
|
# prompts = [p + opt.lora_trigger_word for p in prompts]
|
|
|
|
# go
|
|
start = time.time()
|
|
|
|
pbar = tqdm(range(args.batch_count), leave=False)
|
|
if args.batch_count == 1:
|
|
pbar.disable=True
|
|
|
|
state.job_count = args.batch_count
|
|
|
|
for batch in pbar:
|
|
state.job_no = batch + 1
|
|
if state.skipped:
|
|
state.skipped = False
|
|
|
|
if state.interrupted:
|
|
break
|
|
|
|
state.job = f"Batch {batch+1} out of {args.batch_count}"
|
|
ddim_sampler.noise_gen.manual_seed(args.seed + batch if args.seed != -1 else -1)
|
|
# sample
|
|
samples = sample_text2video(model, args.prompt, args.n_prompt, 1, 1,# todo:add batch size support
|
|
sample_type='ddim', sampler=ddim_sampler,
|
|
ddim_steps=args.steps, eta=args.eta,
|
|
cfg_scale=args.cfg_scale,
|
|
decode_frame_bs=1,
|
|
ddp=False, show_denoising_progress=False,
|
|
num_frames=args.frames
|
|
)
|
|
# save
|
|
if batch > 0:
|
|
outdir_current = os.path.join(get_outdir(), f"{init_timestring}_{batch}")
|
|
print(f'text2video finished, saving frames to {outdir_current}')
|
|
|
|
npz_to_video_grid(samples[0:1,...], # TODO: is this the reason only 1 second is saved?
|
|
os.path.join(outdir_current, f"vid.mp4"),
|
|
fps=video_args.fps)
|
|
if video_args.add_soundtrack != 'None':
|
|
add_soundtrack(video_args.ffmpeg_location, video_args.fps, os.path.join(outdir_current, f"vid.mp4"), 0, -1, None, video_args.add_soundtrack, video_args.soundtrack_path, video_args.ffmpeg_crf, video_args.ffmpeg_preset)
|
|
print(f't2v complete, result saved at {outdir_current}')
|
|
|
|
mp4 = open(outdir_current + os.path.sep + f"vid.mp4", 'rb').read()
|
|
dataurl = "data:video/mp4;base64," + b64encode(mp4).decode()
|
|
t2v_helpers_args.i1_store_t2v = f'<p style=\"font-weight:bold;margin-bottom:0em\">text2video extension for auto1111 — version 1.1b </p><video controls loop><source src="{dataurl}" type="video/mp4"></video>'
|
|
print("Finish sampling!")
|
|
print(f"Run time = {(time.time() - start):.2f} seconds")
|
|
pbar.close()
|
|
return [dataurl]
|