MuseV/musev/models/controlnet.py

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