Upload external/TripoSG/triposg/models/transformers/__init__.py with huggingface_hub

external/TripoSG/triposg/models/transformers/__init__.py

@@ -0,0 +1,61 @@
+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)