400 lines
14 KiB
Python
400 lines
14 KiB
Python
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
|
|
import warnings
|
|
import os
|
|
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from MuseVdiffusers.models.modeling_utils import ModelMixin
|
|
import PIL
|
|
from einops import rearrange, repeat
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.init as init
|
|
from MuseVdiffusers.models.controlnet import ControlNetModel
|
|
from MuseVdiffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
|
|
from MuseVdiffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
|
|
from MuseVdiffusers.utils.torch_utils import is_compiled_module
|
|
|
|
|
|
class ControlnetPredictor(object):
|
|
def __init__(self, controlnet_model_path: str, *args, **kwargs):
|
|
"""Controlnet 推断函数，用于提取 controlnet backbone的emb，避免训练时重复抽取
|
|
Controlnet inference predictor, used to extract the emb of the controlnet backbone to avoid repeated extraction during training
|
|
Args:
|
|
controlnet_model_path (str): controlnet 模型路径. controlnet model path.
|
|
"""
|
|
super(ControlnetPredictor, self).__init__(*args, **kwargs)
|
|
self.controlnet = ControlNetModel.from_pretrained(
|
|
controlnet_model_path,
|
|
)
|
|
|
|
def prepare_image(
|
|
self,
|
|
image, # b c t h w
|
|
width,
|
|
height,
|
|
batch_size,
|
|
num_images_per_prompt,
|
|
device,
|
|
dtype,
|
|
do_classifier_free_guidance=False,
|
|
guess_mode=False,
|
|
):
|
|
if height is None:
|
|
height = image.shape[-2]
|
|
if width is None:
|
|
width = image.shape[-1]
|
|
width, height = (
|
|
x - x % self.control_image_processor.vae_scale_factor
|
|
for x in (width, height)
|
|
)
|
|
image = rearrange(image, "b c t h w-> (b t) c h w")
|
|
image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
|
|
|
|
image = (
|
|
torch.nn.functional.interpolate(
|
|
image,
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
),
|
|
)
|
|
|
|
do_normalize = self.control_image_processor.config.do_normalize
|
|
if image.min() < 0:
|
|
warnings.warn(
|
|
"Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
|
|
f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
|
|
FutureWarning,
|
|
)
|
|
do_normalize = False
|
|
|
|
if do_normalize:
|
|
image = self.control_image_processor.normalize(image)
|
|
|
|
image_batch_size = image.shape[0]
|
|
|
|
if image_batch_size == 1:
|
|
repeat_by = batch_size
|
|
else:
|
|
# image batch size is the same as prompt batch size
|
|
repeat_by = num_images_per_prompt
|
|
|
|
image = image.repeat_interleave(repeat_by, dim=0)
|
|
|
|
image = image.to(device=device, dtype=dtype)
|
|
|
|
if do_classifier_free_guidance and not guess_mode:
|
|
image = torch.cat([image] * 2)
|
|
|
|
return image
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
batch_size: int,
|
|
device: str,
|
|
dtype: torch.dtype,
|
|
timesteps: List[float],
|
|
i: int,
|
|
scheduler: KarrasDiffusionSchedulers,
|
|
prompt_embeds: torch.Tensor,
|
|
do_classifier_free_guidance: bool = False,
|
|
# 2b co t ho wo
|
|
latent_model_input: torch.Tensor = None,
|
|
# b co t ho wo
|
|
latents: torch.Tensor = None,
|
|
# b c t h w
|
|
image: Union[
|
|
torch.FloatTensor,
|
|
PIL.Image.Image,
|
|
np.ndarray,
|
|
List[torch.FloatTensor],
|
|
List[PIL.Image.Image],
|
|
List[np.ndarray],
|
|
] = None,
|
|
# b c t(1) hi wi
|
|
controlnet_condition_frames: Optional[torch.FloatTensor] = None,
|
|
# b c t ho wo
|
|
controlnet_latents: Union[torch.FloatTensor, np.ndarray] = None,
|
|
# b c t(1) ho wo
|
|
controlnet_condition_latents: Optional[torch.FloatTensor] = None,
|
|
height: Optional[int] = None,
|
|
width: Optional[int] = None,
|
|
num_videos_per_prompt: Optional[int] = 1,
|
|
return_dict: bool = True,
|
|
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
|
|
guess_mode: bool = False,
|
|
control_guidance_start: Union[float, List[float]] = 0.0,
|
|
control_guidance_end: Union[float, List[float]] = 1.0,
|
|
latent_index: torch.LongTensor = None,
|
|
vision_condition_latent_index: torch.LongTensor = None,
|
|
**kwargs,
|
|
):
|
|
assert (
|
|
image is None and controlnet_latents is None
|
|
), "should set one of image and controlnet_latents"
|
|
|
|
controlnet = (
|
|
self.controlnet._orig_mod
|
|
if is_compiled_module(self.controlnet)
|
|
else self.controlnet
|
|
)
|
|
|
|
# align format for control guidance
|
|
if not isinstance(control_guidance_start, list) and isinstance(
|
|
control_guidance_end, list
|
|
):
|
|
control_guidance_start = len(control_guidance_end) * [
|
|
control_guidance_start
|
|
]
|
|
elif not isinstance(control_guidance_end, list) and isinstance(
|
|
control_guidance_start, list
|
|
):
|
|
control_guidance_end = len(control_guidance_start) * [control_guidance_end]
|
|
elif not isinstance(control_guidance_start, list) and not isinstance(
|
|
control_guidance_end, list
|
|
):
|
|
mult = (
|
|
len(controlnet.nets)
|
|
if isinstance(controlnet, MultiControlNetModel)
|
|
else 1
|
|
)
|
|
control_guidance_start, control_guidance_end = mult * [
|
|
control_guidance_start
|
|
], mult * [control_guidance_end]
|
|
|
|
if isinstance(controlnet, MultiControlNetModel) and isinstance(
|
|
controlnet_conditioning_scale, float
|
|
):
|
|
controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(
|
|
controlnet.nets
|
|
)
|
|
|
|
global_pool_conditions = (
|
|
controlnet.config.global_pool_conditions
|
|
if isinstance(controlnet, ControlNetModel)
|
|
else controlnet.nets[0].config.global_pool_conditions
|
|
)
|
|
guess_mode = guess_mode or global_pool_conditions
|
|
|
|
# 4. Prepare image
|
|
if isinstance(controlnet, ControlNetModel):
|
|
if (
|
|
controlnet_latents is not None
|
|
and controlnet_condition_latents is not None
|
|
):
|
|
if isinstance(controlnet_latents, np.ndarray):
|
|
controlnet_latents = torch.from_numpy(controlnet_latents)
|
|
if isinstance(controlnet_condition_latents, np.ndarray):
|
|
controlnet_condition_latents = torch.from_numpy(
|
|
controlnet_condition_latents
|
|
)
|
|
# TODO：使用index进行concat
|
|
controlnet_latents = torch.concat(
|
|
[controlnet_condition_latents, controlnet_latents], dim=2
|
|
)
|
|
if not guess_mode and do_classifier_free_guidance:
|
|
controlnet_latents = torch.concat([controlnet_latents] * 2, dim=0)
|
|
controlnet_latents = rearrange(
|
|
controlnet_latents, "b c t h w->(b t) c h w"
|
|
)
|
|
controlnet_latents = controlnet_latents.to(device=device, dtype=dtype)
|
|
else:
|
|
# TODO：使用index进行concat
|
|
# TODO: concat with index
|
|
if controlnet_condition_frames is not None:
|
|
if isinstance(controlnet_condition_frames, np.ndarray):
|
|
image = np.concatenate(
|
|
[controlnet_condition_frames, image], axis=2
|
|
)
|
|
image = self.prepare_image(
|
|
image=image,
|
|
width=width,
|
|
height=height,
|
|
batch_size=batch_size * num_videos_per_prompt,
|
|
num_images_per_prompt=num_videos_per_prompt,
|
|
device=device,
|
|
dtype=controlnet.dtype,
|
|
do_classifier_free_guidance=do_classifier_free_guidance,
|
|
guess_mode=guess_mode,
|
|
)
|
|
height, width = image.shape[-2:]
|
|
elif isinstance(controlnet, MultiControlNetModel):
|
|
images = []
|
|
# TODO: 支持直接使用controlnet_latent而不是frames
|
|
# TODO: support using controlnet_latent directly instead of frames
|
|
if controlnet_latents is not None:
|
|
raise NotImplementedError
|
|
else:
|
|
for i, image_ in enumerate(image):
|
|
if controlnet_condition_frames is not None and isinstance(
|
|
controlnet_condition_frames, list
|
|
):
|
|
if isinstance(controlnet_condition_frames[i], np.ndarray):
|
|
image_ = np.concatenate(
|
|
[controlnet_condition_frames[i], image_], axis=2
|
|
)
|
|
image_ = self.prepare_image(
|
|
image=image_,
|
|
width=width,
|
|
height=height,
|
|
batch_size=batch_size * num_videos_per_prompt,
|
|
num_images_per_prompt=num_videos_per_prompt,
|
|
device=device,
|
|
dtype=controlnet.dtype,
|
|
do_classifier_free_guidance=do_classifier_free_guidance,
|
|
guess_mode=guess_mode,
|
|
)
|
|
|
|
images.append(image_)
|
|
|
|
image = images
|
|
height, width = image[0].shape[-2:]
|
|
else:
|
|
assert False
|
|
|
|
# 7.1 Create tensor stating which controlnets to keep
|
|
controlnet_keep = []
|
|
for i in range(len(timesteps)):
|
|
keeps = [
|
|
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
|
|
for s, e in zip(control_guidance_start, control_guidance_end)
|
|
]
|
|
controlnet_keep.append(
|
|
keeps[0] if isinstance(controlnet, ControlNetModel) else keeps
|
|
)
|
|
|
|
t = timesteps[i]
|
|
|
|
# controlnet(s) inference
|
|
if guess_mode and do_classifier_free_guidance:
|
|
# Infer ControlNet only for the conditional batch.
|
|
control_model_input = latents
|
|
control_model_input = scheduler.scale_model_input(control_model_input, t)
|
|
controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
|
|
else:
|
|
control_model_input = latent_model_input
|
|
controlnet_prompt_embeds = prompt_embeds
|
|
if isinstance(controlnet_keep[i], list):
|
|
cond_scale = [
|
|
c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])
|
|
]
|
|
else:
|
|
cond_scale = controlnet_conditioning_scale * controlnet_keep[i]
|
|
control_model_input_reshape = rearrange(
|
|
control_model_input, "b c t h w -> (b t) c h w"
|
|
)
|
|
encoder_hidden_states_repeat = repeat(
|
|
controlnet_prompt_embeds,
|
|
"b n q->(b t) n q",
|
|
t=control_model_input.shape[2],
|
|
)
|
|
|
|
down_block_res_samples, mid_block_res_sample = self.controlnet(
|
|
control_model_input_reshape,
|
|
t,
|
|
encoder_hidden_states_repeat,
|
|
controlnet_cond=image,
|
|
controlnet_cond_latents=controlnet_latents,
|
|
conditioning_scale=cond_scale,
|
|
guess_mode=guess_mode,
|
|
return_dict=False,
|
|
)
|
|
|
|
return down_block_res_samples, mid_block_res_sample
|
|
|
|
|
|
class InflatedConv3d(nn.Conv2d):
|
|
def forward(self, x):
|
|
video_length = x.shape[2]
|
|
|
|
x = rearrange(x, "b c f h w -> (b f) c h w")
|
|
x = super().forward(x)
|
|
x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
|
|
|
|
return x
|
|
|
|
|
|
def zero_module(module):
|
|
# Zero out the parameters of a module and return it.
|
|
for p in module.parameters():
|
|
p.detach().zero_()
|
|
return module
|
|
|
|
|
|
class PoseGuider(ModelMixin):
|
|
def __init__(
|
|
self,
|
|
conditioning_embedding_channels: int,
|
|
conditioning_channels: int = 3,
|
|
block_out_channels: Tuple[int] = (16, 32, 64, 128),
|
|
):
|
|
super().__init__()
|
|
self.conv_in = InflatedConv3d(
|
|
conditioning_channels, block_out_channels[0], kernel_size=3, padding=1
|
|
)
|
|
|
|
self.blocks = nn.ModuleList([])
|
|
|
|
for i in range(len(block_out_channels) - 1):
|
|
channel_in = block_out_channels[i]
|
|
channel_out = block_out_channels[i + 1]
|
|
self.blocks.append(
|
|
InflatedConv3d(channel_in, channel_in, kernel_size=3, padding=1)
|
|
)
|
|
self.blocks.append(
|
|
InflatedConv3d(
|
|
channel_in, channel_out, kernel_size=3, padding=1, stride=2
|
|
)
|
|
)
|
|
|
|
self.conv_out = zero_module(
|
|
InflatedConv3d(
|
|
block_out_channels[-1],
|
|
conditioning_embedding_channels,
|
|
kernel_size=3,
|
|
padding=1,
|
|
)
|
|
)
|
|
|
|
def forward(self, conditioning):
|
|
embedding = self.conv_in(conditioning)
|
|
embedding = F.silu(embedding)
|
|
|
|
for block in self.blocks:
|
|
embedding = block(embedding)
|
|
embedding = F.silu(embedding)
|
|
|
|
embedding = self.conv_out(embedding)
|
|
|
|
return embedding
|
|
|
|
@classmethod
|
|
def from_pretrained(
|
|
cls,
|
|
pretrained_model_path,
|
|
conditioning_embedding_channels: int,
|
|
conditioning_channels: int = 3,
|
|
block_out_channels: Tuple[int] = (16, 32, 64, 128),
|
|
):
|
|
if not os.path.exists(pretrained_model_path):
|
|
print(f"There is no model file in {pretrained_model_path}")
|
|
print(
|
|
f"loaded PoseGuider's pretrained weights from {pretrained_model_path} ..."
|
|
)
|
|
|
|
state_dict = torch.load(pretrained_model_path, map_location="cpu")
|
|
model = PoseGuider(
|
|
conditioning_embedding_channels=conditioning_embedding_channels,
|
|
conditioning_channels=conditioning_channels,
|
|
block_out_channels=block_out_channels,
|
|
)
|
|
|
|
m, u = model.load_state_dict(state_dict, strict=False)
|
|
# print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
|
|
params = [p.numel() for n, p in model.named_parameters()]
|
|
print(f"### PoseGuider's Parameters: {sum(params) / 1e6} M")
|
|
|
|
return model
|