62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
from typing import Callable, Optional
|
|
|
|
from .triposg_transformer import TripoSGDiTModel
|
|
|
|
|
|
def default_set_attn_proc_func(
|
|
name: str,
|
|
hidden_size: int,
|
|
cross_attention_dim: Optional[int],
|
|
ori_attn_proc: object,
|
|
) -> object:
|
|
return ori_attn_proc
|
|
|
|
|
|
def set_transformer_attn_processor(
|
|
transformer: TripoSGDiTModel,
|
|
set_self_attn_proc_func: Callable = default_set_attn_proc_func,
|
|
set_cross_attn_1_proc_func: Callable = default_set_attn_proc_func,
|
|
set_cross_attn_2_proc_func: Callable = default_set_attn_proc_func,
|
|
set_self_attn_module_names: Optional[list[str]] = None,
|
|
set_cross_attn_1_module_names: Optional[list[str]] = None,
|
|
set_cross_attn_2_module_names: Optional[list[str]] = None,
|
|
) -> None:
|
|
do_set_processor = lambda name, module_names: (
|
|
any([name.startswith(module_name) for module_name in module_names])
|
|
if module_names is not None
|
|
else True
|
|
) # prefix match
|
|
|
|
attn_procs = {}
|
|
for name, attn_processor in transformer.attn_processors.items():
|
|
hidden_size = transformer.config.width
|
|
if name.endswith("attn1.processor"):
|
|
# self attention
|
|
attn_procs[name] = (
|
|
set_self_attn_proc_func(name, hidden_size, None, attn_processor)
|
|
if do_set_processor(name, set_self_attn_module_names)
|
|
else attn_processor
|
|
)
|
|
elif name.endswith("attn2.processor"):
|
|
# cross attention
|
|
cross_attention_dim = transformer.config.cross_attention_dim
|
|
attn_procs[name] = (
|
|
set_cross_attn_1_proc_func(
|
|
name, hidden_size, cross_attention_dim, attn_processor
|
|
)
|
|
if do_set_processor(name, set_cross_attn_1_module_names)
|
|
else attn_processor
|
|
)
|
|
elif name.endswith("attn2_2.processor"):
|
|
# cross attention 2
|
|
cross_attention_dim = transformer.config.cross_attention_2_dim
|
|
attn_procs[name] = (
|
|
set_cross_attn_2_proc_func(
|
|
name, hidden_size, cross_attention_dim, attn_processor
|
|
)
|
|
if do_set_processor(name, set_cross_attn_2_module_names)
|
|
else attn_processor
|
|
)
|
|
|
|
transformer.set_attn_processor(attn_procs)
|