custum_3d_diffusion/custum_modules/unifield_processor.py

from types import FunctionType
from typing import Any, Dict, List
from diffusers import UNet2DConditionModel
import torch
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, ImageProjection
from diffusers.models.attention_processor import Attention, AttnProcessor, AttnProcessor2_0, XFormersAttnProcessor
from dataclasses import dataclass, field
from diffusers.loaders import IPAdapterMixin
from custum_3d_diffusion.custum_modules.attention_processors import add_extra_processor, switch_extra_processor, add_multiview_processor, switch_multiview_processor, add_switch, change_switch

@dataclass
class AttnConfig:
    """        
    * CrossAttention: Attention module (inherits knowledge), LoRA module (achieves fine-tuning), IPAdapter module (achieves conceptual control).
    * SelfAttention: Attention module (inherits knowledge), LoRA module (achieves fine-tuning), Reference Attention module (achieves pixel-level control).
    * Multiview Attention module: Multiview Attention module (achieves multi-view consistency).
    * Cross Modality Attention module: Cross Modality Attention module (achieves multi-modality consistency).
    
    For setups:
        train_xxx_lr is implemented in the U-Net architecture.
        enable_xxx_lora is implemented in the U-Net architecture.
        enable_xxx_ip is implemented in the processor and U-Net architecture.
        enable_xxx_ref_proj_in is implemented in the processor.
    """
    latent_size: int = 64
    
    train_lr: float = 0
    # for cross attention
    # 0 learning rate for not training
    train_cross_attn_lr: float = 0
    train_cross_attn_lora_lr: float = 0       
    train_cross_attn_ip_lr: float = 0      # 0 for not trained
    init_cross_attn_lora: bool = False
    enable_cross_attn_lora: bool = False
    init_cross_attn_ip: bool = False
    enable_cross_attn_ip: bool = False
    cross_attn_lora_rank: int = 64        # 0 for not enabled
    cross_attn_lora_only_kv: bool = False
    ipadapter_pretrained_name: str = "h94/IP-Adapter"
    ipadapter_subfolder_name: str = "models"
    ipadapter_weight_name: str = "ip-adapter-plus_sd15.safetensors"
    ipadapter_effect_on: str = "all"    # all, first

    # for self attention
    train_self_attn_lr: float = 0
    train_self_attn_lora_lr: float = 0
    init_self_attn_lora: bool = False
    enable_self_attn_lora: bool = False
    self_attn_lora_rank: int = 64
    self_attn_lora_only_kv: bool = False

    train_self_attn_ref_lr: float = 0
    train_ref_unet_lr: float = 0
    init_self_attn_ref: bool = False
    enable_self_attn_ref: bool = False      
    self_attn_ref_other_model_name: str = ""
    self_attn_ref_position: str = "attn1"
    self_attn_ref_pixel_wise_crosspond: bool = False    # enable pixel_wise_crosspond in refattn
    self_attn_ref_chain_pos: str = "parralle"           # before or parralle or after
    self_attn_ref_effect_on: str = "all"                # all or first, for _crosspond attn
    self_attn_ref_zero_init: bool = True
    use_simple3d_attn: bool = False

    # for multiview attention
    init_multiview_attn: bool = False
    enable_multiview_attn: bool = False
    multiview_attn_position: str = "attn1"
    multiview_chain_pose: str = "parralle"             # before or parralle or after
    num_modalities: int = 1
    use_mv_joint_attn: bool = False
    
    # for unet
    init_unet_path: str = "stable-diffusion-v1-5/stable-diffusion-v1-5"
    init_num_cls_label: int = 0                         # for initialize
    cls_labels: List[int] = field(default_factory=lambda: [])
    cls_label_type: str = "embedding"
    cat_condition: bool = False                         # cat condition to input

class Configurable:
    attn_config: AttnConfig

    def set_config(self, attn_config: AttnConfig):
        raise NotImplementedError()
    
    def update_config(self, attn_config: AttnConfig):
        self.attn_config = attn_config
    
    def do_set_config(self, attn_config: AttnConfig):
        self.set_config(attn_config)
        for name, module in self.named_modules():
            if isinstance(module, Configurable):
                if hasattr(module, "do_set_config"):
                    module.do_set_config(attn_config)
                else:
                    print(f"Warning: {name} has no attribute do_set_config, but is an instance of Configurable")
                    module.attn_config = attn_config

    def do_update_config(self, attn_config: AttnConfig):
        self.update_config(attn_config)
        for name, module in self.named_modules():
            if isinstance(module, Configurable):
                if hasattr(module, "do_update_config"):
                    module.do_update_config(attn_config)
                else:
                    print(f"Warning: {name} has no attribute do_update_config, but is an instance of Configurable")
                    module.attn_config = attn_config

from diffusers import ModelMixin  # Must import ModelMixin for CompiledUNet
class UnifieldWrappedUNet(UNet2DConditionModel):
    forward_hook: FunctionType

    def forward(self, *args, **kwargs):
        if hasattr(self, 'forward_hook'):
            return self.forward_hook(super().forward, *args, **kwargs)
        return super().forward(*args, **kwargs)


class ConfigurableUNet2DConditionModel(Configurable, IPAdapterMixin):
    unet: UNet2DConditionModel

    cls_embedding_param_dict = {}
    cross_attn_lora_param_dict = {}
    self_attn_lora_param_dict = {}
    cross_attn_param_dict = {}
    self_attn_param_dict = {}
    ipadapter_param_dict = {}
    ref_attn_param_dict = {}
    ref_unet_param_dict = {}
    multiview_attn_param_dict = {}
    other_param_dict = {}
    
    rev_param_name_mapping = {}

    class_labels = []
    def set_class_labels(self, class_labels: torch.Tensor):
        if self.attn_config.init_num_cls_label != 0:
            self.class_labels = class_labels.to(self.unet.device).long()

    def __init__(self, init_config: AttnConfig, weight_dtype) -> None:
        super().__init__()
        self.weight_dtype = weight_dtype
        self.set_config(init_config)

    def enable_xformers_memory_efficient_attention(self):
        self.unet.enable_xformers_memory_efficient_attention
        def recursive_add_processors(name: str, module: torch.nn.Module):
            for sub_name, child in module.named_children():
                recursive_add_processors(f"{name}.{sub_name}", child)

            if isinstance(module, Attention):
                if hasattr(module, 'xformers_not_supported'):
                    return
                old_processor = module.get_processor()
                if isinstance(old_processor, (AttnProcessor, AttnProcessor2_0)):
                    module.set_use_memory_efficient_attention_xformers(True)

        for name, module in self.unet.named_children():
            recursive_add_processors(name, module)

    def __getattr__(self, name: str) -> Any:
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.unet, name)
    
    # --- for IPAdapterMixin
    
    def register_modules(self, **kwargs):
        for name, module in kwargs.items():
            # set models
            setattr(self, name, module)

    def register_to_config(self, **kwargs):
        pass

    def unload_ip_adapter(self):
        raise NotImplementedError()

    # --- for Configurable

    def get_refunet(self):
        if self.attn_config.self_attn_ref_other_model_name == "self":
            return self.unet
        else:
            return self.unet.ref_unet

    def set_config(self, attn_config: AttnConfig):
        self.attn_config = attn_config

        unet_type = UnifieldWrappedUNet
        # class_embed_type = "projection" for 'camera'
        # class_embed_type = None for 'embedding'
        unet_kwargs = {}
        if attn_config.init_num_cls_label > 0:
            if attn_config.cls_label_type == "embedding":
                unet_kwargs = {
                    "num_class_embeds": attn_config.init_num_cls_label, 
                    "device_map": None, 
                    "low_cpu_mem_usage": False,
                    "class_embed_type": None,
                }
            else:
                raise ValueError(f"cls_label_type {attn_config.cls_label_type} is not supported")
        
        self.unet: UnifieldWrappedUNet = unet_type.from_pretrained(
            attn_config.init_unet_path, subfolder="unet", torch_dtype=self.weight_dtype, 
            **unet_kwargs
        )
        assert isinstance(self.unet, UnifieldWrappedUNet)
        self.unet.forward_hook = self.unet_forward_hook

        if self.attn_config.cat_condition:
            # double in_channels
            if self.unet.config.in_channels != 8:
                self.unet.register_to_config(in_channels=self.unet.config.in_channels * 2)
                # repeate unet.conv_in weight twice
                doubled_conv_in = torch.nn.Conv2d(self.unet.conv_in.in_channels * 2, self.unet.conv_in.out_channels, self.unet.conv_in.kernel_size, self.unet.conv_in.stride, self.unet.conv_in.padding)
                doubled_conv_in.weight.data = torch.cat([self.unet.conv_in.weight.data, torch.zeros_like(self.unet.conv_in.weight.data)], dim=1)
                doubled_conv_in.bias.data = self.unet.conv_in.bias.data
                self.unet.conv_in = doubled_conv_in
        
        used_param_ids = set()
        
        if attn_config.init_cross_attn_lora:
            # setup lora
            from peft import LoraConfig
            from peft.utils import get_peft_model_state_dict
            if attn_config.cross_attn_lora_only_kv:
                target_modules=["attn2.to_k", "attn2.to_v"]
            else:
                target_modules=["attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0"]
            lora_config: LoraConfig = LoraConfig(
                r=attn_config.cross_attn_lora_rank,
                lora_alpha=attn_config.cross_attn_lora_rank,
                init_lora_weights="gaussian",
                target_modules=target_modules,
            )
            adapter_name="cross_attn_lora"
            self.unet.add_adapter(lora_config, adapter_name=adapter_name)
            # update cross_attn_lora_param_dict
            self.cross_attn_lora_param_dict = {id(param): param for name, param in self.unet.named_parameters() if adapter_name in name and id(param) not in used_param_ids}
            used_param_ids.update(self.cross_attn_lora_param_dict.keys())

        if attn_config.init_self_attn_lora:
            # setup lora
            from peft import LoraConfig
            if attn_config.self_attn_lora_only_kv:
                target_modules=["attn1.to_k", "attn1.to_v"]
            else:
                target_modules=["attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0"]
            lora_config: LoraConfig = LoraConfig(
                r=attn_config.self_attn_lora_rank,
                lora_alpha=attn_config.self_attn_lora_rank,
                init_lora_weights="gaussian",
                target_modules=target_modules,
            )
            adapter_name="self_attn_lora"
            self.unet.add_adapter(lora_config, adapter_name=adapter_name)
            # update cross_self_lora_param_dict
            self.self_attn_lora_param_dict = {id(param): param for name, param in self.unet.named_parameters() if adapter_name in name and id(param) not in used_param_ids}
            used_param_ids.update(self.self_attn_lora_param_dict.keys())

        if attn_config.init_num_cls_label != 0:
            self.cls_embedding_param_dict = {id(param): param for param in self.unet.class_embedding.parameters()}
            used_param_ids.update(self.cls_embedding_param_dict.keys())
            self.set_class_labels(torch.tensor(attn_config.cls_labels).long())
        
        if attn_config.init_cross_attn_ip:
            self.image_encoder = None
            # setup ipadapter
            self.load_ip_adapter(
                attn_config.ipadapter_pretrained_name,
                subfolder=attn_config.ipadapter_subfolder_name,
                weight_name=attn_config.ipadapter_weight_name
            )
            # warp ip_adapter_attn_proc with switch
            from diffusers.models.attention_processor import IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0
            add_switch(self.unet, module_filter=lambda x: isinstance(x, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0)), switch_dict_fn=lambda x: {"ipadapter": x, "default": XFormersAttnProcessor()}, switch_name="ipadapter_switch", enabled_proc="ipadapter")
            # update ipadapter_param_dict
            # weights are in attention processors and unet.encoder_hid_proj
            self.ipadapter_param_dict = {id(param): param for param in self.unet.encoder_hid_proj.parameters() if id(param) not in used_param_ids}
            used_param_ids.update(self.ipadapter_param_dict.keys())
            print("DEBUG: ipadapter_param_dict len in encoder_hid_proj", len(self.ipadapter_param_dict))
            for name, processor in self.unet.attn_processors.items():
                if hasattr(processor, "to_k_ip"):
                    self.ipadapter_param_dict.update({id(param): param for param in processor.parameters()})
            print(f"DEBUG: ipadapter_param_dict len in all", len(self.ipadapter_param_dict))

        ref_unet = None
        if attn_config.init_self_attn_ref:
            # setup reference attention processor
            if attn_config.self_attn_ref_other_model_name == "self":
                raise NotImplementedError("self reference is not fully implemented")
            else:
                ref_unet: UNet2DConditionModel = UNet2DConditionModel.from_pretrained(
                    attn_config.self_attn_ref_other_model_name, subfolder="unet", torch_dtype=self.unet.dtype
                )
                ref_unet.to(self.unet.device)
                if self.attn_config.train_ref_unet_lr == 0:
                    ref_unet.eval()
                    ref_unet.requires_grad_(False)
                else:
                    ref_unet.train()

                add_extra_processor(
                    model=ref_unet, 
                    enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"), 
                    mode='extract',
                    with_proj_in=False,
                    pixel_wise_crosspond=False,
                )
                # NOTE: Here require cross_attention_dim in two unet's self attention should be the same
                processor_dict = add_extra_processor(
                    model=self.unet,
                    enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"),
                    mode='inject',
                    with_proj_in=False,
                    pixel_wise_crosspond=attn_config.self_attn_ref_pixel_wise_crosspond,
                    crosspond_effect_on=attn_config.self_attn_ref_effect_on,
                    crosspond_chain_pos=attn_config.self_attn_ref_chain_pos,
                    simple_3d=attn_config.use_simple3d_attn,
                )
                self.ref_unet_param_dict = {id(param): param for name, param in ref_unet.named_parameters() if id(param) not in used_param_ids and (attn_config.self_attn_ref_position in name)}
                if attn_config.self_attn_ref_chain_pos != "after":
                    # pop untrainable paramters
                    for name, param in ref_unet.named_parameters():
                        if id(param) in self.ref_unet_param_dict and ('up_blocks.3.attentions.2.transformer_blocks.0.' in name):
                            self.ref_unet_param_dict.pop(id(param))
                used_param_ids.update(self.ref_unet_param_dict.keys())
            # update ref_attn_param_dict
            self.ref_attn_param_dict = {id(param): param for name, param in processor_dict.named_parameters() if id(param) not in used_param_ids}
            used_param_ids.update(self.ref_attn_param_dict.keys())

        if attn_config.init_multiview_attn:
            processor_dict = add_multiview_processor(
                model = self.unet, 
                enable_filter = lambda name: name.endswith(f"{attn_config.multiview_attn_position}.processor"),
                num_modalities = attn_config.num_modalities,    
                base_img_size = attn_config.latent_size,      
                chain_pos = attn_config.multiview_chain_pose,
            )
            # update multiview_attn_param_dict
            self.multiview_attn_param_dict = {id(param): param for name, param in processor_dict.named_parameters() if id(param) not in used_param_ids}
            used_param_ids.update(self.multiview_attn_param_dict.keys())
        
        # initialize cross_attn_param_dict parameters
        self.cross_attn_param_dict = {id(param): param for name, param in self.unet.named_parameters() if "attn2" in name and id(param) not in used_param_ids}
        used_param_ids.update(self.cross_attn_param_dict.keys())
        
        # initialize self_attn_param_dict parameters
        self.self_attn_param_dict = {id(param): param for name, param in self.unet.named_parameters() if "attn1" in name and id(param) not in used_param_ids}
        used_param_ids.update(self.self_attn_param_dict.keys())
        
        # initialize other_param_dict parameters
        self.other_param_dict = {id(param): param for name, param in self.unet.named_parameters() if id(param) not in used_param_ids}
        
        if ref_unet is not None:
            self.unet.ref_unet = ref_unet
            
        self.rev_param_name_mapping = {id(param): name for name, param in self.unet.named_parameters()}
        
        self.update_config(attn_config, force_update=True)
        return self.unet
    
    _attn_keys_to_update = ["enable_cross_attn_lora", "enable_cross_attn_ip", "enable_self_attn_lora", "enable_self_attn_ref", "enable_multiview_attn", "cls_labels"]
    
    def update_config(self, attn_config: AttnConfig, force_update=False):
        assert isinstance(self.unet, UNet2DConditionModel), "unet must be an instance of UNet2DConditionModel"

        need_to_update = False
        # update cls_labels
        for key in self._attn_keys_to_update:
            if getattr(self.attn_config, key) != getattr(attn_config, key):
                need_to_update = True
                break
        if not force_update and not need_to_update:
            return

        self.set_class_labels(torch.tensor(attn_config.cls_labels).long())
        
        # setup loras
        if self.attn_config.init_cross_attn_lora or self.attn_config.init_self_attn_lora:
            if attn_config.enable_cross_attn_lora or attn_config.enable_self_attn_lora:
                cross_attn_lora_weight = 1. if attn_config.enable_cross_attn_lora > 0 else 0
                self_attn_lora_weight = 1. if attn_config.enable_self_attn_lora > 0 else 0
                self.unet.set_adapters(["cross_attn_lora", "self_attn_lora"], weights=[cross_attn_lora_weight, self_attn_lora_weight])
            else:
                self.unet.disable_adapters()

        # setup ipadapter
        if self.attn_config.init_cross_attn_ip:
            if attn_config.enable_cross_attn_ip:
                change_switch(self.unet, "ipadapter_switch", "ipadapter")
            else:
                change_switch(self.unet, "ipadapter_switch", "default")
            
        # setup reference attention processor
        if self.attn_config.init_self_attn_ref:
            if attn_config.enable_self_attn_ref:
                switch_extra_processor(self.unet, enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"))
            else:
                switch_extra_processor(self.unet, enable_filter=lambda name: False)
        
        # setup multiview attention processor
        if self.attn_config.init_multiview_attn:
            if attn_config.enable_multiview_attn:
                switch_multiview_processor(self.unet, enable_filter=lambda name: name.endswith(f"{attn_config.multiview_attn_position}.processor"))
            else:
                switch_multiview_processor(self.unet, enable_filter=lambda name: False)
        
        # update cls_labels
        for key in self._attn_keys_to_update:
            setattr(self.attn_config, key, getattr(attn_config, key))

    def unet_forward_hook(self, raw_forward, sample: torch.FloatTensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, *args, cross_attention_kwargs=None, condition_latents=None, class_labels=None, noisy_condition_input=False, cond_pixels_clip=None, **kwargs):
        if class_labels is None and len(self.class_labels) > 0:
            class_labels = self.class_labels.repeat(sample.shape[0] // self.class_labels.shape[0]).to(sample.device)
        elif self.attn_config.init_num_cls_label != 0:
            assert class_labels is not None, "class_labels should be passed if self.class_labels is empty and self.attn_config.init_num_cls_label is not 0"
        if class_labels is not None:
            if self.attn_config.cls_label_type == "embedding":
                pass
            else:
                raise ValueError(f"cls_label_type {self.attn_config.cls_label_type} is not supported")
        if self.attn_config.init_self_attn_ref and self.attn_config.enable_self_attn_ref:
            # NOTE: extra step, extract condition
            ref_dict = {}
            ref_unet = self.get_refunet().to(sample.device)
            assert condition_latents is not None
            if self.attn_config.self_attn_ref_other_model_name == "self":
                raise NotImplementedError()
            else:
                with torch.no_grad():
                    cond_encoder_hidden_states = encoder_hidden_states.reshape(condition_latents.shape[0], -1, *encoder_hidden_states.shape[1:])[:, 0]
                    if timestep.dim() == 0:
                        cond_timestep = timestep
                    else:
                        cond_timestep = timestep.reshape(condition_latents.shape[0], -1)[:, 0]
                ref_unet(condition_latents, cond_timestep, cond_encoder_hidden_states,  cross_attention_kwargs=dict(ref_dict=ref_dict))
            # NOTE: extra step, inject condition
            # Predict the noise residual and compute loss
            if cross_attention_kwargs is None:
                cross_attention_kwargs = {}
            cross_attention_kwargs.update(ref_dict=ref_dict, mode='inject')
        elif condition_latents is not None:
            if not hasattr(self, 'condition_latents_raised'):
                print("Warning! condition_latents is not None, but self_attn_ref is not enabled! This warning will only be raised once.")
                self.condition_latents_raised = True
        
        if self.attn_config.init_cross_attn_ip:
            raise NotImplementedError()
        
        if self.attn_config.cat_condition:
            assert condition_latents is not None
            B = condition_latents.shape[0]
            cat_latents = condition_latents.reshape(B, 1, *condition_latents.shape[1:]).repeat(1, sample.shape[0] // B, 1, 1, 1).reshape(*sample.shape)
            sample = torch.cat([sample, cat_latents], dim=1)
            
        return raw_forward(sample, timestep, encoder_hidden_states, *args, cross_attention_kwargs=cross_attention_kwargs, class_labels=class_labels, **kwargs)