Spaces:

HTBris
/

zerogpu-sdxlTrainer

Running on Zero

zerogpu-sdxlTrainer / toolkit /models /v2 /diffusion_models /hidream.py

Deploy Ostris AI-Toolkit ZeroGPU Trainer (Illustrious XL / SDXL / FLUX.1)

855e1f5 3 days ago

541 Bytes

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"]

	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"]