19 lines
541 B
Python
19 lines
541 B
Python
from diffusers.models import (
|
|
HiDreamImageTransformer2DModel as DiffusersHiDreamImageTransformer2DModel,
|
|
)
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
class HiDreamImageTransformer2DModel(
|
|
DiffusersHiDreamImageTransformer2DModel, OstrisModelMixin
|
|
):
|
|
"""The diffusers HiDream DiT (hidream_e1; the base hidream arch uses the
|
|
vendored copy in the hidream extension)."""
|
|
|
|
aitk_subfolder = "transformer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["double_stream_blocks", "single_stream_blocks"]
|