import os
import warnings
from typing import Optional
import torch
from diffusers import DDIMScheduler, TextToVideoSDPipeline
from einops import rearrange
from torch import Tensor
from torch.nn.functional import interpolate
from tqdm import trange
import random
from MotionDirector_train import export_to_video, handle_memory_attention, load_primary_models, unet_and_text_g_c, freeze_models
from utils.lora_handler import LoraHandler
from utils.ddim_utils import ddim_inversion
from utils.lora import extract_lora_child_module
import imageio
def initialize_pipeline(
model: str,
device: str = "cuda",
xformers: bool = True,
sdp: bool = True,
lora_path: str = "",
lora_rank: int = 32,
lora_scale: float = 1.0,
):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(model)
# Freeze any necessary models
freeze_models([vae, text_encoder, unet])
# Enable xformers if available
handle_memory_attention(xformers, sdp, unet)
lora_manager_temporal = LoraHandler(
version="cloneofsimo",
use_unet_lora=True,
use_text_lora=False,
save_for_webui=False,
only_for_webui=False,
unet_replace_modules=["TransformerTemporalModel"],
text_encoder_replace_modules=None,
lora_bias=None
)
unet_lora_params, unet_negation = lora_manager_temporal.add_lora_to_model(
True, unet, lora_manager_temporal.unet_replace_modules, 0, lora_path, r=lora_rank, scale=lora_scale)
unet.eval()
text_encoder.eval()
unet_and_text_g_c(unet, text_encoder, False, False)
pipe = TextToVideoSDPipeline.from_pretrained(
pretrained_model_name_or_path=model,
scheduler=scheduler,
tokenizer=tokenizer,
text_encoder=text_encoder.to(device=device, dtype=torch.half),
vae=vae.to(device=device, dtype=torch.half),
unet=unet.to(device=device, dtype=torch.half),
)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
return pipe
def inverse_video(pipe, latents, num_steps):
ddim_inv_scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
ddim_inv_scheduler.set_timesteps(num_steps)
ddim_inv_latent = ddim_inversion(
pipe, ddim_inv_scheduler, video_latent=latents.to(pipe.device),
num_inv_steps=num_steps, prompt="")[-1]
return ddim_inv_latent
def prepare_input_latents(
pipe: TextToVideoSDPipeline,
batch_size: int,
num_frames: int,
height: int,
width: int,
latents_path:str,
model_select: str,
random_seed: int,
):
# initialize with random gaussian noise
scale = pipe.vae_scale_factor
shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
if random_seed > 1000:
torch.manual_seed(random_seed)
else:
random_seed = random.randint(100, 10000000)
torch.manual_seed(random_seed)
print(f"random_seed: {random_seed}")
if '1-' in model_select:
noise_prior = 0.3
elif '2-' in model_select:
noise_prior = 0.5
elif '3-' in model_select:
noise_prior = 0.
else:
noise_prior = 0.
if noise_prior > 0.:
cached_latents = torch.load(latents_path)
if 'inversion_noise' not in cached_latents:
latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
else:
latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
if latents.shape[0] != batch_size:
latents = latents.repeat(batch_size, 1, 1, 1, 1)
if latents.shape != shape:
latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
noise = torch.randn_like(latents, dtype=torch.half)
latents_base = noise
latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
else:
latents = torch.randn(shape, dtype=torch.half)
latents_base = latents
return latents, latents_base, random_seed
class MotionDirector():
def __init__(self):
self.version = "0.0.0"
self.foundation_model_path = "./zeroscope_v2_576w/"
self.lora_path = "./MotionDirector_pretrained/dolly_zoom_(hitchcockian_zoom)/checkpoint-default/temporal/lora"
with torch.autocast("cuda", dtype=torch.half):
self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path, lora_scale=1)
def reload_lora(self, lora_path):
if lora_path != self.lora_path:
self.lora_path = lora_path
with torch.autocast("cuda", dtype=torch.half):
self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path)
def __call__(self, model_select, text_pormpt, neg_text_pormpt, random_seed, steps, guidance_scale, baseline_select):
model_select = str(model_select)
out_name = f"./outputs/inference"
out_name += f"{text_pormpt}".replace(' ', '_').replace(',', '').replace('.', '')
model_select_type = model_select.split('--')[1].strip()
model_select_type = model_select_type.lower().replace(' ', '_')
lora_path = f"./MotionDirector_pretrained/{model_select_type}/checkpoint-default/temporal/lora"
self.reload_lora(lora_path)
latents_folder = f"./MotionDirector_pretrained/{model_select_type}/cached_latents"
latents_path = f"{latents_folder}/{random.choice(os.listdir(latents_folder))}"
assert os.path.exists(lora_path)
device = "cuda"
self.pipe.enable_vae_slicing()
with torch.autocast(device, dtype=torch.half):
# prepare input latents
with torch.no_grad():
init_latents, init_latents_base, random_seed = prepare_input_latents(
pipe=self.pipe,
batch_size=1,
num_frames=16,
height=384,
width=384,
latents_path=latents_path,
model_select=model_select,
random_seed=random_seed
)
video_frames = self.pipe(
prompt=text_pormpt,
negative_prompt=neg_text_pormpt,
width=384,
height=384,
num_frames=16,
num_inference_steps=steps,
guidance_scale=guidance_scale,
latents=init_latents
).frames
out_file = f"{out_name}_{random_seed}.mp4"
os.makedirs(os.path.dirname(out_file), exist_ok=True)
export_to_video(video_frames, out_file, 8)
if baseline_select:
with torch.autocast("cuda", dtype=torch.half):
loras = extract_lora_child_module(self.pipe.unet, target_replace_module=["TransformerTemporalModel"])
for lora_i in loras:
lora_i.scale = 0.
# self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path,
# lora_scale=0.)
with torch.no_grad():
video_frames = self.pipe(
prompt=text_pormpt,
negative_prompt=neg_text_pormpt,
width=384,
height=384,
num_frames=16,
num_inference_steps=steps,
guidance_scale=guidance_scale,
latents=init_latents_base,
).frames
out_file_baseline = f"{out_name}_{random_seed}_baseline.mp4"
os.makedirs(os.path.dirname(out_file_baseline), exist_ok=True)
export_to_video(video_frames, out_file_baseline, 8)
# with torch.autocast("cuda", dtype=torch.half):
# self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path,
# lora_scale=1.)
loras = extract_lora_child_module(self.pipe.unet,
target_replace_module=["TransformerTemporalModel"])
for lora_i in loras:
lora_i.scale = 1.
else:
out_file_baseline = None
return [out_file, out_file_baseline]