from typing import Any, Dict, List, Optional, Tuple, Union from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from diffusers.configuration_utils import register_to_config from diffusers.utils import BaseOutput from models.controlnet_sdv import ControlNetSDVModel, zero_module # from unimatch.unimatch.geometry import flow_warp from models.softsplat import softsplat # from models.hourglass.dense_motion import DenseMotionNetwork import models.cmp.models as cmp_models import models.cmp.utils as cmp_utils import yaml import os import torchvision.transforms as transforms class ArgObj(object): def __init__(self): pass class CMP_demo(nn.Module): def __init__(self, configfn, load_iter): super().__init__() args = ArgObj() with open(configfn) as f: config = yaml.full_load(f) for k, v in config.items(): setattr(args, k, v) setattr(args, 'load_iter', load_iter) setattr(args, 'exp_path', os.path.dirname(configfn)) self.model = cmp_models.__dict__[args.model['arch']](args.model, dist_model=False) self.model.load_state("{}/checkpoints".format(args.exp_path), args.load_iter, False) self.model.switch_to('eval') self.data_mean = args.data['data_mean'] self.data_div = args.data['data_div'] self.img_transform = transforms.Compose([ transforms.Normalize(self.data_mean, self.data_div)]) self.args = args self.fuser = cmp_utils.Fuser(args.model['module']['nbins'], args.model['module']['fmax']) torch.cuda.synchronize() def run(self, image, sparse, mask): dtype = image.dtype image = image * 2 - 1 self.model.set_input(image.float(), torch.cat([sparse, mask], dim=1).float(), None) cmp_output = self.model.model(self.model.image_input, self.model.sparse_input) flow = self.fuser.convert_flow(cmp_output) if flow.shape[2] != self.model.image_input.shape[2]: flow = nn.functional.interpolate( flow, size=self.model.image_input.shape[2:4], mode="bilinear", align_corners=True) return flow.to(dtype) # [b, 2, h, w] # tensor_dict = self.model.eval(ret_loss=False) # flow = tensor_dict['flow_tensors'][0].cpu().numpy().squeeze().transpose(1,2,0) # return flow class FlowControlNetConditioningEmbeddingSVD(nn.Module): def __init__( self, conditioning_embedding_channels: int, conditioning_channels: int = 3, block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), ): super().__init__() self.conv_in = nn.Conv2d(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(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) self.conv_out = zero_module( nn.Conv2d(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 class FlowControlNetFirstFrameEncoderLayer(nn.Module): def __init__( self, c_in, c_out, is_downsample=False ): super().__init__() self.conv_in = nn.Conv2d(c_in, c_out, kernel_size=3, padding=1, stride=2 if is_downsample else 1) def forward(self, feature): ''' feature: [b, c, h, w] ''' embedding = self.conv_in(feature) embedding = F.silu(embedding) return embedding class FlowControlNetFirstFrameEncoder(nn.Module): def __init__( self, c_in=320, channels=[320, 640, 1280], downsamples=[True, True, True], use_zeroconv=True ): super().__init__() self.encoders = nn.ModuleList([]) self.zeroconvs = nn.ModuleList([]) for channel, downsample in zip(channels, downsamples): self.encoders.append(FlowControlNetFirstFrameEncoderLayer(c_in, channel, is_downsample=downsample)) self.zeroconvs.append(zero_module(nn.Conv2d(channel, channel, kernel_size=1)) if use_zeroconv else nn.Identity()) c_in = channel def forward(self, first_frame): feature = first_frame deep_features = [] for encoder, zeroconv in zip(self.encoders, self.zeroconvs): feature = encoder(feature) # print(feature.shape) deep_features.append(zeroconv(feature)) return deep_features @dataclass class FlowControlNetOutput(BaseOutput): """ The output of [`FlowControlNetOutput`]. Args: down_block_res_samples (`tuple[torch.Tensor]`): A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be used to condition the original UNet's downsampling activations. mid_down_block_re_sample (`torch.Tensor`): The activation of the midde block (the lowest sample resolution). Each tensor should be of shape `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`. Output can be used to condition the original UNet's middle block activation. """ down_block_res_samples: Tuple[torch.Tensor] mid_block_res_sample: torch.Tensor controlnet_flow: torch.Tensor cmp_output: torch.Tensor class FlowControlNet(ControlNetSDVModel): _supports_gradient_checkpointing = True @register_to_config def __init__( self, sample_size: Optional[int] = None, in_channels: int = 8, out_channels: int = 4, down_block_types: Tuple[str] = ( "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal", ), up_block_types: Tuple[str] = ( "UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", ), block_out_channels: Tuple[int] = (320, 640, 1280, 1280), addition_time_embed_dim: int = 256, projection_class_embeddings_input_dim: int = 768, layers_per_block: Union[int, Tuple[int]] = 2, cross_attention_dim: Union[int, Tuple[int]] = 1024, transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20), num_frames: int = 25, conditioning_channels: int = 3, conditioning_embedding_out_channels : Optional[Tuple[int, ...]] = (16, 32, 96, 256), ): super().__init__() self.flow_encoder = FlowControlNetFirstFrameEncoder() # time_embed_dim = block_out_channels[0] * 4 # blocks_time_embed_dim = time_embed_dim self.controlnet_cond_embedding = FlowControlNetConditioningEmbeddingSVD( conditioning_embedding_channels=block_out_channels[0], block_out_channels=conditioning_embedding_out_channels, conditioning_channels=conditioning_channels, ) def get_warped_frames(self, first_frame, flows): ''' video_frame: [b, c, w, h] flows: [b, t-1, c, w, h] ''' dtype = first_frame.dtype warped_frames = [] for i in range(flows.shape[1]): warped_frame = softsplat(tenIn=first_frame.float(), tenFlow=flows[:, i].float(), tenMetric=None, strMode='avg').to(dtype) # [b, c, w, h] warped_frames.append(warped_frame.unsqueeze(1)) # [b, 1, c, w, h] warped_frames = torch.cat(warped_frames, dim=1) # [b, t-1, c, w, h] return warped_frames def get_cmp_flow(self, frames, sparse_optical_flow, mask): ''' frames: [b, 13, 3, 384, 384] (0, 1) tensor sparse_optical_flow: [b, 13, 2, 384, 384] (-384, 384) tensor mask: [b, 13, 2, 384, 384] {0, 1} tensor ''' b, t, c, h, w = frames.shape assert h == 384 and w == 384 frames = frames.flatten(0, 1) # [b*13, 3, 256, 256] sparse_optical_flow = sparse_optical_flow.flatten(0, 1) # [b*13, 2, 256, 256] mask = mask.flatten(0, 1) # [b*13, 2, 256, 256] cmp_flow, cmp_output = self.run(frames, sparse_optical_flow, mask) # [b*13, 2, 256, 256] # cmp_flow = self.run(frames.float(), sparse_optical_flow.float(), mask.float()) # [b*13, 2, 256, 256] cmp_flow = cmp_flow.reshape(b, t, 2, h, w) return cmp_flow, cmp_output # return cmp_flow.to(dtype=dtype) def run(self, image, sparse, mask): image = image * 2 - 1 cmp_output = self.cmp_model(image, torch.cat([sparse, mask], dim=1)) flow = self.fuser.convert_flow(cmp_output) if flow.shape[2] != image.shape[2]: flow = nn.functional.interpolate( flow, size=image.shape[2:4], mode="bilinear", align_corners=True) return flow, cmp_output # [b, 2, h, w] def forward( self, sample: torch.FloatTensor, timestep: Union[torch.Tensor, float, int], encoder_hidden_states: torch.Tensor, added_time_ids: torch.Tensor, controlnet_cond: torch.FloatTensor = None, # [b, 3, h, w] controlnet_flow: torch.FloatTensor = None, # [b, 13, 2, h, w] # controlnet_mask: torch.FloatTensor = None, # [b, 13, 2, h, w] # pixel_values_384: torch.FloatTensor = None, # sparse_optical_flow_384: torch.FloatTensor = None, # mask_384: torch.FloatTensor = None, image_only_indicator: Optional[torch.Tensor] = None, return_dict: bool = True, guess_mode: bool = False, conditioning_scale: float = 1.0, ) -> Union[FlowControlNetOutput, Tuple]: # print(sample.shape) # print(controlnet_cond.shape) # print(controlnet_flow.shape) # assert False # 1. time timesteps = timestep if not torch.is_tensor(timesteps): # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can # This would be a good case for the `match` statement (Python 3.10+) is_mps = sample.device.type == "mps" if isinstance(timestep, float): dtype = torch.float32 if is_mps else torch.float64 else: dtype = torch.int32 if is_mps else torch.int64 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) elif len(timesteps.shape) == 0: timesteps = timesteps[None].to(sample.device) # broadcast to batch dimension in a way that's compatible with ONNX/Core ML batch_size, num_frames = sample.shape[:2] timesteps = timesteps.expand(batch_size) t_emb = self.time_proj(timesteps) # `Timesteps` does not contain any weights and will always return f32 tensors # but time_embedding might actually be running in fp16. so we need to cast here. # there might be better ways to encapsulate this. t_emb = t_emb.to(dtype=sample.dtype) emb = self.time_embedding(t_emb) time_embeds = self.add_time_proj(added_time_ids.flatten()) time_embeds = time_embeds.reshape((batch_size, -1)) time_embeds = time_embeds.to(emb.dtype) aug_emb = self.add_embedding(time_embeds) emb = emb + aug_emb # Flatten the batch and frames dimensions # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width] sample = sample.flatten(0, 1) # Repeat the embeddings num_video_frames times # emb: [batch, channels] -> [batch * frames, channels] emb = emb.repeat_interleave(num_frames, dim=0) # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels] encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0) # hourglass_output = self.hourglass_forward( # controlnet_cond, controlnet_sparse_flow, controlnet_mask, controlnet_init_flow) # [b, l, 3+2+2, h, w] # controlnet_flow = controlnet_init_flow + hourglass_output # 2. pre-process sample = self.conv_in(sample) # [b*l, 320, h//8, w//8] # print(controlnet_cond.shape) # controlnet cond if controlnet_cond != None: # embed 成 64*64，和latent一个shape controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) # [b, 320, h//8, w//8] # sample = sample + controlnet_cond # print(controlnet_cond.shape) # assert False controlnet_cond_features = [controlnet_cond] + self.flow_encoder(controlnet_cond) # [4] # print(controlnet_cond.shape) ''' torch.Size([2, 320, 32, 32]) torch.Size([2, 320, 16, 16]) torch.Size([2, 640, 8, 8]) torch.Size([2, 1280, 4, 4]) ''' # for x in controlnet_cond_features: # print(x.shape) # assert False scales = [8, 16, 32, 64] scale_flows = {} fb, fl, fc, fh, fw = controlnet_flow.shape # print(controlnet_flow.shape) for scale in scales: scaled_flow = F.interpolate(controlnet_flow.reshape(-1, fc, fh, fw), scale_factor=1/scale) scaled_flow = scaled_flow.reshape(fb, fl, fc, fh // scale, fw // scale) / scale scale_flows[scale] = scaled_flow # for k in scale_flows.keys(): # print(scale_flows[k].shape) # assert False warped_cond_features = [] for cond_feature in controlnet_cond_features: cb, cc, ch, cw = cond_feature.shape # print(cond_feature.shape) warped_cond_feature = self.get_warped_frames(cond_feature, scale_flows[fh // ch]) warped_cond_feature = torch.cat([cond_feature.unsqueeze(1), warped_cond_feature], dim=1) # [b, c, h, w] wb, wl, wc, wh, ww = warped_cond_feature.shape # print(warped_cond_feature.shape) warped_cond_features.append(warped_cond_feature.reshape(wb * wl, wc, wh, ww)) # for x in warped_cond_features: # print(x.shape) # assert False ''' torch.Size([28, 320, 32, 32]) torch.Size([28, 320, 16, 16]) torch.Size([28, 640, 8, 8]) torch.Size([28, 1280, 4, 4]) ''' image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device) count = 0 length = len(warped_cond_features) # print(sample.shape) # print(warped_cond_features[0].shape) # add the warped feature in the first scale sample = sample + warped_cond_features[count] count += 1 down_block_res_samples = (sample,) # print(sample.shape) for downsample_block in self.down_blocks: if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: sample, res_samples = downsample_block( hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states, image_only_indicator=image_only_indicator, ) else: sample, res_samples = downsample_block( hidden_states=sample, temb=emb, image_only_indicator=image_only_indicator, ) # print(sample.shape) # print(warped_cond_features[min(count, length - 1)].shape) sample = sample + warped_cond_features[min(count, length - 1)] count += 1 down_block_res_samples += res_samples # print(len(res_samples)) # for i in range(len(res_samples)): # print(res_samples[i].shape) # [28, 320, 32, 32] # [28, 320, 32, 32] # [28, 320, 16, 16] # [28, 640, 16, 16] # [28, 640, 16, 16] # [28, 640, 8, 8] # [28, 1280, 8, 8] # [28, 1280, 8, 8] # [28, 1280, 4, 4] # [28, 1280, 4, 4] # [28, 1280, 4, 4] # print(sample.shape) # print(warped_cond_features[-1].shape) # add the warped feature in the last scale sample = sample + warped_cond_features[-1] # 4. mid sample = self.mid_block( hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states, image_only_indicator=image_only_indicator, ) # [b*l, 1280, h // 64, w // 64] # print(sample.shape) # assert False controlnet_down_block_res_samples = () for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): down_block_res_sample = controlnet_block(down_block_res_sample) controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,) down_block_res_samples = controlnet_down_block_res_samples mid_block_res_sample = self.controlnet_mid_block(sample) # 6. scaling down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples] mid_block_res_sample = mid_block_res_sample * conditioning_scale # for sample in down_block_res_samples: # print(sample.shape) # print(mid_block_res_sample.shape) # assert False if not return_dict: return (down_block_res_samples, mid_block_res_sample, controlnet_flow, None) return FlowControlNetOutput( down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample, controlnet_flow=controlnet_flow, cmp_output=None )