block.py

"""Entry point for the Krea 2 reference-image edit modular blocks (loaded via `trust_remote_code`).

Importing this module registers the Krea 2 classes into the `diffusers` namespace / modular mapping so that,
on a stock `diffusers` install (which does not ship Krea 2), component loading and `init_pipeline` resolve:

- `Krea2Transformer2DModel` — referenced as `["diffusers", "Krea2Transformer2DModel"]` in the base repo's
  `model_index.json`;
- `Krea2ModularPipeline` — the model-specific pipeline class (`init_pipeline` maps `model_name="krea2"` to it;
  without this it would fall back to the generic `ModularPipeline`, which lacks the Krea 2 properties such as
  `requires_unconditional_embeds`, `vae_scale_factor`, `num_channels_latents`).
"""

import diffusers
from diffusers.modular_pipelines.modular_pipeline import (
    MODULAR_PIPELINE_MAPPING,
    _create_default_map_fn,
)

from .transformer_krea2 import Krea2Transformer2DModel  # noqa: F401  (registers into diffusers namespace)
from .modular_pipeline import Krea2ModularPipeline
from .modular_blocks_krea2_edit import Krea2EditBlocks


diffusers.Krea2ModularPipeline = Krea2ModularPipeline
MODULAR_PIPELINE_MAPPING.setdefault("krea2", _create_default_map_fn("Krea2ModularPipeline"))


__all__ = ["Krea2EditBlocks", "Krea2Transformer2DModel", "Krea2ModularPipeline"]