329 lines
18 KiB
Python
329 lines
18 KiB
Python
from pathlib import Path
|
|
from typing import Any, Optional, Union, Callable
|
|
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
from diffusers import DDPMScheduler, DiffusionPipeline, AutoencoderKL, DDIMScheduler
|
|
from diffusers.utils.import_utils import is_xformers_available
|
|
from einops import rearrange, repeat
|
|
|
|
from transformers import CLIPTextModel, CLIPTokenizer
|
|
from ..utils.video_utils import ResultProcessor, save_videos_grid, video_naming
|
|
|
|
from . import pl_module_params_controlnet
|
|
|
|
from .diffusers_conditional.models.controlnet.controlnet import ControlNetModel
|
|
from .diffusers_conditional.models.controlnet.unet_3d_condition import UNet3DConditionModel
|
|
from .diffusers_conditional.models.controlnet.pipeline_text_to_video_w_controlnet_synth import TextToVideoSDPipeline
|
|
|
|
from .diffusers_conditional.models.controlnet.processor import set_use_memory_efficient_attention_xformers
|
|
from .diffusers_conditional.models.controlnet.mask_generator import MaskGenerator
|
|
|
|
import warnings
|
|
# from warnings import warn
|
|
from ..utils.iimage import IImage
|
|
from ..utils.object_loader import instantiate_object
|
|
from ..utils.object_loader import get_class
|
|
|
|
|
|
class VideoLDM(pl.LightningModule):
|
|
|
|
def __init__(self,
|
|
inference_params: pl_module_params_controlnet.InferenceParams,
|
|
opt_params: pl_module_params_controlnet.OptimizerParams = None,
|
|
unet_params: pl_module_params_controlnet.UNetParams = None,
|
|
):
|
|
super().__init__()
|
|
|
|
self.inference_generator = torch.Generator(device=self.device)
|
|
|
|
self.opt_params = opt_params
|
|
self.unet_params = unet_params
|
|
|
|
print(f"Base pipeline from: {unet_params.pipeline_repo}")
|
|
print(f"Pipeline class {unet_params.pipeline_class}")
|
|
# load entire pipeline (unet, vq, text encoder,..)
|
|
state_dict_control_model = None
|
|
state_dict_fusion = None
|
|
state_dict_base_model = None
|
|
|
|
if len(opt_params.load_trained_controlnet_from_ckpt) > 0:
|
|
state_dict_ckpt = torch.load(opt_params.load_trained_controlnet_from_ckpt, map_location=torch.device("cpu"))
|
|
state_dict_ckpt = state_dict_ckpt["state_dict"]
|
|
state_dict_control_model = dict(filter(lambda x: x[0].startswith("unet"), state_dict_ckpt.items()))
|
|
state_dict_control_model = {k.split("unet.")[1]: v for (k, v) in state_dict_control_model.items()}
|
|
|
|
state_dict_fusion = dict(filter(lambda x: "cross_attention_merger" in x[0], state_dict_ckpt.items()))
|
|
state_dict_fusion = {k.split("base_model.")[1]: v for (k, v) in state_dict_fusion.items()}
|
|
del state_dict_ckpt
|
|
|
|
state_dict_proj = None
|
|
state_dict_ckpt = None
|
|
|
|
if hasattr(unet_params, "use_resampler") and unet_params.use_resampler:
|
|
num_queries = unet_params.num_frames if unet_params.num_frames > 1 else None
|
|
if unet_params.use_image_tokens_ctrl:
|
|
num_queries = unet_params.num_control_input_frames
|
|
assert unet_params.frame_expansion == "none"
|
|
image_encoder = self.unet_params.image_encoder
|
|
embedding_dim = image_encoder.embedding_dim
|
|
|
|
resampler = instantiate_object(self.unet_params.resampler_cls, video_length=num_queries, embedding_dim=embedding_dim, input_tokens=image_encoder.num_tokens, num_layers=self.unet_params.resampler_merging_layers, aggregation=self.unet_params.aggregation)
|
|
|
|
state_dict_proj = None
|
|
|
|
self.resampler = resampler
|
|
self.image_encoder = image_encoder
|
|
|
|
|
|
noise_scheduler = DDPMScheduler.from_pretrained(self.unet_params.pipeline_repo, subfolder="scheduler")
|
|
tokenizer = CLIPTokenizer.from_pretrained(self.unet_params.pipeline_repo, subfolder="tokenizer")
|
|
text_encoder = CLIPTextModel.from_pretrained(self.unet_params.pipeline_repo, subfolder="text_encoder")
|
|
vae = AutoencoderKL.from_pretrained(self.unet_params.pipeline_repo, subfolder="vae")
|
|
base_model = UNet3DConditionModel.from_pretrained(self.unet_params.pipeline_repo, subfolder="unet", low_cpu_mem_usage=False, device_map=None, merging_mode=self.unet_params.merging_mode_base, use_image_embedding=unet_params.use_resampler and unet_params.use_image_tokens_main, use_fps_conditioning=self.opt_params.use_fps_conditioning, unet_params=unet_params)
|
|
|
|
if state_dict_base_model is not None:
|
|
miss, unex = base_model.load_state_dict(state_dict_base_model, strict=False)
|
|
assert len(unex) == 0
|
|
if len(miss) > 0:
|
|
warnings.warn(f"Missing keys when loading base_mode:{miss}")
|
|
del state_dict_base_model
|
|
if state_dict_fusion is not None:
|
|
miss, unex = base_model.load_state_dict(state_dict_fusion, strict=False)
|
|
assert len(unex) == 0
|
|
del state_dict_fusion
|
|
|
|
print("PIPE LOADING DONE")
|
|
self.noise_scheduler = noise_scheduler
|
|
self.tokenizer = tokenizer
|
|
self.text_encoder = text_encoder
|
|
self.vae = vae
|
|
|
|
self.unet = ControlNetModel.from_unet(
|
|
unet=base_model,
|
|
conditioning_embedding_out_channels=unet_params.conditioning_embedding_out_channels,
|
|
downsample_controlnet_cond=unet_params.downsample_controlnet_cond,
|
|
num_frames=unet_params.num_frames if (unet_params.frame_expansion != "none" or self.unet_params.use_controlnet_mask) else unet_params.num_control_input_frames,
|
|
num_frame_conditioning=unet_params.num_control_input_frames,
|
|
frame_expansion=unet_params.frame_expansion,
|
|
pre_transformer_in_cond=unet_params.pre_transformer_in_cond,
|
|
num_tranformers=unet_params.num_tranformers,
|
|
vae=AutoencoderKL.from_pretrained(self.unet_params.pipeline_repo, subfolder="vae"),
|
|
zero_conv_mode=unet_params.zero_conv_mode,
|
|
merging_mode=unet_params.merging_mode,
|
|
condition_encoder=unet_params.condition_encoder,
|
|
use_controlnet_mask=unet_params.use_controlnet_mask,
|
|
use_image_embedding=unet_params.use_resampler and unet_params.use_image_tokens_ctrl,
|
|
unet_params=unet_params,
|
|
use_image_encoder_normalization=unet_params.use_image_encoder_normalization,
|
|
)
|
|
if state_dict_control_model is not None:
|
|
miss, unex = self.unet.load_state_dict(
|
|
state_dict_control_model, strict=False)
|
|
if len(miss) > 0:
|
|
print("WARNING: Loading checkpoint for controlnet misses states")
|
|
print(miss)
|
|
|
|
if unet_params.frame_expansion == "none":
|
|
attention_params = self.unet_params.attention_mask_params
|
|
assert not attention_params.temporal_self_attention_only_on_conditioning and not attention_params.spatial_attend_on_condition_frames and not attention_params.temp_attend_on_neighborhood_of_condition_frames
|
|
|
|
self.mask_generator = MaskGenerator(
|
|
self.unet_params.attention_mask_params, num_frame_conditioning=self.unet_params.num_control_input_frames, num_frames=self.unet_params.num_frames)
|
|
self.mask_generator_base = MaskGenerator(
|
|
self.unet_params.attention_mask_params_base, num_frame_conditioning=self.unet_params.num_control_input_frames, num_frames=self.unet_params.num_frames)
|
|
|
|
if state_dict_proj is not None and unet_params.use_image_tokens_main:
|
|
if unet_params.use_image_tokens_main:
|
|
missing, unexpected = base_model.load_state_dict(
|
|
state_dict_proj, strict=False)
|
|
elif unet_params.use_image_tokens_ctrl:
|
|
missing, unexpected = unet.load_state_dict(
|
|
state_dict_proj, strict=False)
|
|
assert len(unexpected) == 0, f"Unexpected entries {unexpected}"
|
|
print(f"Missing keys state proj = {missing}")
|
|
del state_dict_proj
|
|
|
|
base_model.requires_grad_(False)
|
|
self.base_model = base_model
|
|
self.unet.requires_grad_(False)
|
|
self.text_encoder.requires_grad_(False)
|
|
self.vae.requires_grad_(False)
|
|
|
|
layers_config = opt_params.layers_config
|
|
layers_config.set_requires_grad(self)
|
|
|
|
print("CUSTOM XFORMERS ATTENTION USED.")
|
|
if is_xformers_available():
|
|
set_use_memory_efficient_attention_xformers(self.unet, num_frame_conditioning=self.unet_params.num_control_input_frames,
|
|
num_frames=self.unet_params.num_frames,
|
|
attention_mask_params=self.unet_params.attention_mask_params
|
|
)
|
|
set_use_memory_efficient_attention_xformers(self.base_model, num_frame_conditioning=self.unet_params.num_control_input_frames,
|
|
num_frames=self.unet_params.num_frames,
|
|
attention_mask_params=self.unet_params.attention_mask_params_base)
|
|
|
|
if len(inference_params.scheduler_cls) > 0:
|
|
inf_scheduler_class = get_class(inference_params.scheduler_cls)
|
|
else:
|
|
inf_scheduler_class = DDIMScheduler
|
|
|
|
inf_scheduler = inf_scheduler_class.from_pretrained(
|
|
self.unet_params.pipeline_repo, subfolder="scheduler")
|
|
inference_pipeline = TextToVideoSDPipeline(vae=self.vae,
|
|
text_encoder=self.text_encoder,
|
|
tokenizer=self.tokenizer,
|
|
unet=self.base_model,
|
|
controlnet=self.unet,
|
|
scheduler=inf_scheduler
|
|
)
|
|
|
|
inference_pipeline.set_noise_generator(self.opt_params.noise_generator)
|
|
inference_pipeline.enable_vae_slicing()
|
|
|
|
inference_pipeline.set_progress_bar_config(disable=True)
|
|
|
|
self.inference_params = inference_params
|
|
self.inference_pipeline = inference_pipeline
|
|
|
|
self.result_processor = ResultProcessor(fps=self.inference_params.frame_rate, n_frames=self.inference_params.video_length)
|
|
|
|
def on_start(self):
|
|
datamodule = self.trainer._data_connector._datahook_selector.datamodule
|
|
pipe_id_model = self.unet_params.pipeline_repo
|
|
for dataset_key in ["video_dataset", "image_dataset", "predict_dataset"]:
|
|
dataset = getattr(datamodule, dataset_key, None)
|
|
if dataset is not None and hasattr(dataset, "model_id"):
|
|
pipe_id_data = dataset.model_id
|
|
assert pipe_id_model == pipe_id_data, f"Model and Dataloader need the same pipeline path. Found '{pipe_id_model}' and '{dataset_key}.model_id={pipe_id_data}'. Consider setting '--data.{dataset_key}.model_id={pipe_id_data}'"
|
|
self.result_processor.set_logger(self.logger)
|
|
|
|
def on_predict_start(self) -> None:
|
|
self.on_start()
|
|
# pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16")
|
|
# pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
|
# pipe.set_progress_bar_config(disable=True)
|
|
# self.first_stage = pipe.to(self.device)
|
|
|
|
def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
|
|
cfg = self.trainer.predict_cfg
|
|
|
|
result_file_stem = cfg["result_file_stem"]
|
|
storage_fol = Path(cfg['predict_dir'])
|
|
prompts = [cfg["prompt"]]
|
|
|
|
inference_params: pl_module_params_controlnet.InferenceParams = self.inference_params
|
|
conditioning_type = inference_params.conditioning_type
|
|
# n_autoregressive_generations = inference_params.n_autoregressive_generations
|
|
n_autoregressive_generations = cfg["n_autoregressive_generations"]
|
|
mode = inference_params.mode
|
|
start_from_real_input = inference_params.start_from_real_input
|
|
assert isinstance(prompts, list)
|
|
|
|
prompts = n_autoregressive_generations * prompts
|
|
|
|
self.inference_generator.manual_seed(self.inference_params.seed)
|
|
|
|
assert self.unet_params.num_control_input_frames == self.inference_params.video_length//2, f"currently we assume to have an equal size for and second half of the frame interval, e.g. 16 frames, and we condition on 8. Current setup: {self.unet_params.num_frame_conditioning} and {self.inference_params.video_length}"
|
|
|
|
chunks_conditional = []
|
|
batch_size = 1
|
|
shape = (batch_size, self.inference_pipeline.unet.config.in_channels, self.inference_params.video_length,
|
|
self.inference_pipeline.unet.config.sample_size, self.inference_pipeline.unet.config.sample_size)
|
|
for idx, prompt in enumerate(prompts):
|
|
if idx > 0:
|
|
content = sample*2-1
|
|
content_latent = self.vae.encode(content).latent_dist.sample() * self.vae.config.scaling_factor
|
|
content_latent = rearrange(content_latent, "F C W H -> 1 C F W H")
|
|
content_latent = content_latent[:, :, self.unet_params.num_control_input_frames:].detach().clone()
|
|
|
|
if hasattr(self.inference_pipeline, "noise_generator"):
|
|
latents = self.inference_pipeline.noise_generator.sample_noise(shape=shape, device=self.device, dtype=self.dtype, generator=self.inference_generator, content=content_latent if idx > 0 else None)
|
|
else:
|
|
latents = None
|
|
if idx == 0:
|
|
sample = cfg["video"]
|
|
else:
|
|
if inference_params.conditioning_type == "fixed":
|
|
context = chunks_conditional[0][:self.unet_params.num_frame_conditioning]
|
|
context = [context]
|
|
context = [2*sample-1 for sample in context]
|
|
|
|
input_frames_conditioning = torch.cat(context).detach().clone()
|
|
input_frames_conditioning = rearrange(input_frames_conditioning, "F C W H -> 1 F C W H")
|
|
elif inference_params.conditioning_type == "last_chunk":
|
|
input_frames_conditioning = condition_input[:, -self.unet_params.num_frame_conditioning:].detach().clone()
|
|
elif inference_params.conditioning_type == "past":
|
|
context = [sample[:self.unet_params.num_control_input_frames] for sample in chunks_conditional]
|
|
context = [2*sample-1 for sample in context]
|
|
|
|
input_frames_conditioning = torch.cat(context).detach().clone()
|
|
input_frames_conditioning = rearrange(input_frames_conditioning, "F C W H -> 1 F C W H")
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
input_frames = condition_input[:, self.unet_params.num_control_input_frames:].detach().clone()
|
|
|
|
sample = self(prompt, input_frames=input_frames, input_frames_conditioning=input_frames_conditioning, latents=latents)
|
|
|
|
if hasattr(self.inference_pipeline, "reset_noise_generator_state"):
|
|
self.inference_pipeline.reset_noise_generator_state()
|
|
|
|
condition_input = rearrange(sample, "F C W H -> 1 F C W H")
|
|
condition_input = (2*condition_input)-1 # range: [-1,1]
|
|
|
|
# store first 16 frames, then always last 8 of a chunk
|
|
chunks_conditional.append(sample)
|
|
|
|
result_formats = self.inference_params.result_formats
|
|
# result_formats = [gif", "mp4"]
|
|
concat_video = self.inference_params.concat_video
|
|
|
|
def IImage_normalized(x): return IImage(x, vmin=0, vmax=1)
|
|
for result_format in result_formats:
|
|
save_format = result_format.replace("eval_", "")
|
|
|
|
merged_video = None
|
|
for chunk_idx, (prompt, video) in enumerate(zip(prompts, chunks_conditional)):
|
|
if chunk_idx == 0:
|
|
current_video = IImage_normalized(video)
|
|
else:
|
|
current_video = IImage_normalized(video[self.unet_params.num_control_input_frames:])
|
|
|
|
if merged_video is None:
|
|
merged_video = current_video
|
|
else:
|
|
merged_video &= current_video
|
|
|
|
if concat_video:
|
|
filename = video_naming(prompts[0], save_format, batch_idx, 0)
|
|
result_file_video = (storage_fol / filename).absolute().as_posix()
|
|
result_file_video = (Path(result_file_video).parent / (result_file_stem+Path(result_file_video).suffix)).as_posix()
|
|
self.result_processor.save_to_file(video=merged_video.torch(vmin=0, vmax=1), prompt=prompts[0], video_filename=result_file_video, prompt_on_vid=False)
|
|
|
|
def forward(self, prompt, input_frames=None, input_frames_conditioning=None, latents=None):
|
|
call_params = self.inference_params.to_dict()
|
|
print(f"INFERENCE PARAMS = {call_params}")
|
|
call_params["prompt"] = prompt
|
|
|
|
call_params["image"] = input_frames
|
|
call_params["num_frames"] = self.inference_params.video_length
|
|
call_params["return_dict"] = False
|
|
call_params["output_type"] = "pt_t2v"
|
|
call_params["mask_generator"] = self.mask_generator
|
|
call_params["precision"] = "16" if self.trainer.precision.startswith("16") else "32"
|
|
call_params["no_text_condition_control"] = self.opt_params.no_text_condition_control
|
|
call_params["weight_control_sample"] = self.unet_params.weight_control_sample
|
|
call_params["use_controlnet_mask"] = self.unet_params.use_controlnet_mask
|
|
call_params["skip_controlnet_branch"] = self.opt_params.skip_controlnet_branch
|
|
call_params["img_cond_resampler"] = self.resampler if self.unet_params.use_resampler else None
|
|
call_params["img_cond_encoder"] = self.image_encoder if self.unet_params.use_resampler else None
|
|
call_params["input_frames_conditioning"] = input_frames_conditioning
|
|
call_params["cfg_text_image"] = self.unet_params.cfg_text_image
|
|
call_params["use_of"] = self.unet_params.use_of
|
|
if latents is not None:
|
|
call_params["latents"] = latents
|
|
|
|
sample = self.inference_pipeline(generator=self.inference_generator, **call_params)
|
|
return sample
|