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 diffusers.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 diffusers.models.controlnet import ControlNetModel
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
from diffusers.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