import os import json from typing import Any, Dict, Optional from diffusers.models import UNet2DConditionModel import numpy import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint import torch.distributed from PIL import Image from einops import rearrange from typing import Any, Callable, Dict, List, Optional, Union, Tuple import diffusers from diffusers import ( AutoencoderKL, DDPMScheduler, DiffusionPipeline, EulerAncestralDiscreteScheduler, UNet2DConditionModel, ImagePipelineOutput ) from diffusers.image_processor import VaeImageProcessor from diffusers.models.attention_processor import Attention, AttnProcessor, XFormersAttnProcessor, AttnProcessor2_0 from diffusers.utils.import_utils import is_xformers_available from diffusers.utils import deprecate from diffusers.models.transformers.transformer_2d import BasicTransformerBlock def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int): # "feed_forward_chunk_size" can be used to save memory if hidden_states.shape[chunk_dim] % chunk_size != 0: raise ValueError( f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." ) num_chunks = hidden_states.shape[chunk_dim] // chunk_size ff_output = torch.cat( [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)], dim=chunk_dim, ) return ff_output class Basic2p5DTransformerBlock(torch.nn.Module): def __init__(self, transformer: BasicTransformerBlock, layer_name, use_ma=True, use_ra=True) -> None: super().__init__() self.transformer = transformer self.layer_name = layer_name self.use_ma = use_ma self.use_ra = use_ra # multiview attn if self.use_ma: self.attn_multiview = Attention( query_dim=self.dim, heads=self.num_attention_heads, dim_head=self.attention_head_dim, dropout=self.dropout, bias=self.attention_bias, cross_attention_dim=None, upcast_attention=self.attn1.upcast_attention, out_bias=True, ) # ref attn if self.use_ra: self.attn_refview = Attention( query_dim=self.dim, heads=self.num_attention_heads, dim_head=self.attention_head_dim, dropout=self.dropout, bias=self.attention_bias, cross_attention_dim=None, upcast_attention=self.attn1.upcast_attention, out_bias=True, ) def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.transformer, name) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, timestep: Optional[torch.LongTensor] = None, cross_attention_kwargs: Dict[str, Any] = None, class_labels: Optional[torch.LongTensor] = None, added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, ) -> torch.Tensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Self-Attention batch_size = hidden_states.shape[0] cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} num_in_batch = cross_attention_kwargs.pop('num_in_batch', 1) mode = cross_attention_kwargs.pop('mode', None) mva_scale = cross_attention_kwargs.pop('mva_scale', 1.0) ref_scale = cross_attention_kwargs.pop('ref_scale', 1.0) condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None) if self.norm_type == "ada_norm": norm_hidden_states = self.norm1(hidden_states, timestep) elif self.norm_type == "ada_norm_zero": norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype ) elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: norm_hidden_states = self.norm1(hidden_states) elif self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) elif self.norm_type == "ada_norm_single": shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa else: raise ValueError("Incorrect norm used") if self.pos_embed is not None: norm_hidden_states = self.pos_embed(norm_hidden_states) # 1. Prepare GLIGEN inputs cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} gligen_kwargs = cross_attention_kwargs.pop("gligen", None) attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, attention_mask=attention_mask, **cross_attention_kwargs, ) if self.norm_type == "ada_norm_zero": attn_output = gate_msa.unsqueeze(1) * attn_output elif self.norm_type == "ada_norm_single": attn_output = gate_msa * attn_output hidden_states = attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # 1.2 Reference Attention if 'w' in mode: condition_embed_dict[self.layer_name] = rearrange(norm_hidden_states, '(b n) l c -> b (n l) c', n=num_in_batch) # B, (N L), C if 'r' in mode and self.use_ra: condition_embed = condition_embed_dict[self.layer_name].unsqueeze(1).repeat(1,num_in_batch,1,1) # B N L C condition_embed = rearrange(condition_embed, 'b n l c -> (b n) l c') attn_output = self.attn_refview( norm_hidden_states, encoder_hidden_states=condition_embed, attention_mask=None, **cross_attention_kwargs ) ref_scale_timing = ref_scale if isinstance(ref_scale, torch.Tensor): ref_scale_timing = ref_scale.unsqueeze(1).repeat(1, num_in_batch).view(-1) for _ in range(attn_output.ndim - 1): ref_scale_timing = ref_scale_timing.unsqueeze(-1) hidden_states = ref_scale_timing * attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # 1.3 Multiview Attention if num_in_batch > 1 and self.use_ma: multivew_hidden_states = rearrange(norm_hidden_states, '(b n) l c -> b (n l) c', n=num_in_batch) attn_output = self.attn_multiview( multivew_hidden_states, encoder_hidden_states=multivew_hidden_states, **cross_attention_kwargs ) attn_output = rearrange(attn_output, 'b (n l) c -> (b n) l c', n=num_in_batch) hidden_states = mva_scale * attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # 1.2 GLIGEN Control if gligen_kwargs is not None: hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) # 3. Cross-Attention if self.attn2 is not None: if self.norm_type == "ada_norm": norm_hidden_states = self.norm2(hidden_states, timestep) elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: norm_hidden_states = self.norm2(hidden_states) elif self.norm_type == "ada_norm_single": # For PixArt norm2 isn't applied here: # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 norm_hidden_states = hidden_states elif self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"]) else: raise ValueError("Incorrect norm") if self.pos_embed is not None and self.norm_type != "ada_norm_single": norm_hidden_states = self.pos_embed(norm_hidden_states) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, **cross_attention_kwargs, ) hidden_states = attn_output + hidden_states # 4. Feed-forward # i2vgen doesn't have this norm 🤷‍♂️ if self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"]) elif not self.norm_type == "ada_norm_single": norm_hidden_states = self.norm3(hidden_states) if self.norm_type == "ada_norm_zero": norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] if self.norm_type == "ada_norm_single": norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp if self._chunk_size is not None: # "feed_forward_chunk_size" can be used to save memory ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) else: ff_output = self.ff(norm_hidden_states) if self.norm_type == "ada_norm_zero": ff_output = gate_mlp.unsqueeze(1) * ff_output elif self.norm_type == "ada_norm_single": ff_output = gate_mlp * ff_output hidden_states = ff_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) return hidden_states import copy class UNet2p5DConditionModel(torch.nn.Module): def __init__(self, unet: UNet2DConditionModel) -> None: super().__init__() self.unet = unet self.use_ma = True self.use_ra = True self.use_camera_embedding = True self.use_dual_stream = True if self.use_dual_stream: self.unet_dual = copy.deepcopy(unet) self.init_attention(self.unet_dual) self.init_attention(self.unet, use_ma=self.use_ma, use_ra=self.use_ra) self.init_condition() self.init_camera_embedding() @staticmethod def from_pretrained(pretrained_model_name_or_path, **kwargs): torch_dtype = kwargs.pop('torch_dtype', torch.float32) config_path = os.path.join(pretrained_model_name_or_path, 'config.json') unet_ckpt_path = os.path.join(pretrained_model_name_or_path, 'diffusion_pytorch_model.bin') with open(config_path, 'r', encoding='utf-8') as file: config = json.load(file) unet = UNet2DConditionModel(**config) unet = UNet2p5DConditionModel(unet) unet_ckpt = torch.load(unet_ckpt_path, map_location='cpu', weights_only=True) unet.load_state_dict(unet_ckpt, strict=True) unet = unet.to(torch_dtype) return unet def init_condition(self): self.unet.conv_in = torch.nn.Conv2d( 12, self.unet.conv_in.out_channels, kernel_size=self.unet.conv_in.kernel_size, stride=self.unet.conv_in.stride, padding=self.unet.conv_in.padding, dilation=self.unet.conv_in.dilation, groups=self.unet.conv_in.groups, bias=self.unet.conv_in.bias is not None) self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1,77,1024)) self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1,77,1024)) def init_camera_embedding(self): self.max_num_ref_image = 5 self.max_num_gen_image = 12*3+4*2 if self.use_camera_embedding: time_embed_dim = 1280 self.unet.class_embedding = nn.Embedding(self.max_num_ref_image+self.max_num_gen_image, time_embed_dim) def init_attention(self, unet, use_ma=False, use_ra=False): for down_block_i, down_block in enumerate(unet.down_blocks): if hasattr(down_block, "has_cross_attention") and down_block.has_cross_attention: for attn_i, attn in enumerate(down_block.attentions): for transformer_i, transformer in enumerate(attn.transformer_blocks): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'down_{down_block_i}_{attn_i}_{transformer_i}', use_ma, use_ra) if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention: for attn_i, attn in enumerate(unet.mid_block.attentions): for transformer_i, transformer in enumerate(attn.transformer_blocks): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'mid_{attn_i}_{transformer_i}', use_ma, use_ra) for up_block_i, up_block in enumerate(unet.up_blocks): if hasattr(up_block, "has_cross_attention") and up_block.has_cross_attention: for attn_i, attn in enumerate(up_block.attentions): for transformer_i, transformer in enumerate(attn.transformer_blocks): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'up_{up_block_i}_{attn_i}_{transformer_i}', use_ma, use_ra) def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.unet, name) def forward( self, sample, timestep, encoder_hidden_states, *args, down_intrablock_additional_residuals=None, down_block_res_samples=None, mid_block_res_sample=None, **cached_condition, ): B, N_gen, _, H, W = sample.shape assert H == W if self.use_camera_embedding: camera_info_gen = cached_condition['camera_info_gen'] + self.max_num_ref_image camera_info_gen = rearrange(camera_info_gen, 'b n -> (b n)') else: camera_info_gen = None sample = [sample] if 'normal_imgs' in cached_condition: sample.append(cached_condition["normal_imgs"]) if 'position_imgs' in cached_condition: sample.append(cached_condition["position_imgs"]) sample = torch.cat(sample, dim=2) sample = rearrange(sample, 'b n c h w -> (b n) c h w') encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat(1, N_gen, 1, 1) encoder_hidden_states_gen = rearrange(encoder_hidden_states_gen, 'b n l c -> (b n) l c') if self.use_ra: if 'condition_embed_dict' in cached_condition: condition_embed_dict = cached_condition['condition_embed_dict'] else: condition_embed_dict = {} ref_latents = cached_condition['ref_latents'] N_ref = ref_latents.shape[1] if self.use_camera_embedding: camera_info_ref = cached_condition['camera_info_ref'] camera_info_ref = rearrange(camera_info_ref, 'b n -> (b n)') else: camera_info_ref = None ref_latents = rearrange(ref_latents, 'b n c h w -> (b n) c h w') encoder_hidden_states_ref = self.unet.learned_text_clip_ref.unsqueeze(1).repeat(B, N_ref, 1, 1) encoder_hidden_states_ref = rearrange(encoder_hidden_states_ref, 'b n l c -> (b n) l c') noisy_ref_latents = ref_latents timestep_ref = 0 if self.use_dual_stream: unet_ref = self.unet_dual else: unet_ref = self.unet unet_ref( noisy_ref_latents, timestep_ref, encoder_hidden_states=encoder_hidden_states_ref, class_labels=camera_info_ref, # **kwargs return_dict=False, cross_attention_kwargs={ 'mode':'w', 'num_in_batch':N_ref, 'condition_embed_dict':condition_embed_dict}, ) cached_condition['condition_embed_dict'] = condition_embed_dict else: condition_embed_dict = None mva_scale = cached_condition.get('mva_scale', 1.0) ref_scale = cached_condition.get('ref_scale', 1.0) return self.unet( sample, timestep, encoder_hidden_states_gen, *args, class_labels=camera_info_gen, down_intrablock_additional_residuals=[ sample.to(dtype=self.unet.dtype) for sample in down_intrablock_additional_residuals ] if down_intrablock_additional_residuals is not None else None, down_block_additional_residuals=[ sample.to(dtype=self.unet.dtype) for sample in down_block_res_samples ] if down_block_res_samples is not None else None, mid_block_additional_residual=( mid_block_res_sample.to(dtype=self.unet.dtype) if mid_block_res_sample is not None else None ), return_dict=False, cross_attention_kwargs={ 'mode':'r', 'num_in_batch':N_gen, 'condition_embed_dict':condition_embed_dict, 'mva_scale': mva_scale, 'ref_scale': ref_scale, }, )