16 lines
402 B
Python
16 lines
402 B
Python
from diffusers.models import (
|
|
CosmosTransformer3DModel as DiffusersCosmosTransformer3DModel,
|
|
)
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
class CosmosTransformer3DModel(DiffusersCosmosTransformer3DModel, OstrisModelMixin):
|
|
"""Cosmos video/image DiT (anima)."""
|
|
|
|
aitk_subfolder = "transformer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["transformer_blocks"]
|