14 lines
392 B
Python
14 lines
392 B
Python
from diffusers import AuraFlowTransformer2DModel as DiffusersAuraFlowTransformer2DModel
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
class AuraFlowTransformer2DModel(
|
|
DiffusersAuraFlowTransformer2DModel, OstrisModelMixin
|
|
):
|
|
aitk_subfolder = "transformer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["joint_transformer_blocks", "single_transformer_blocks"]
|