# Credits: Built on open-source libraries and papers acknowledged in README.md citations.
"""PyTorch image-to-image translation models.

A library for multi-modal image translation with diffusion bridges,
GANs, and transformer backbones.
"""

__version__ = "0.4.1"

from src.data import PairedImageDataset, UnpairedImageDataset, default_transforms, get_transforms
from src.pipelines import ImageTranslator
from src.losses import GANLoss, PerceptualLoss
from src.metrics import (
    PairedImageMetricEvaluator,
    compute_dists,
    compute_fid,
    compute_lpips,
    compute_psnr,
    compute_ssim,
)
from src.models import I2SBUNet, PatchGANDiscriminator, ResNetGenerator, SIT_CONFIGS, SiTBackbone, UNetGenerator
from src.models import ResnetMaskV1Generator, ResnetMaskV3Generator, NetMatchability
from src.models import BBDMUNet, BDBMUNet, BiBBDMUNet, CDTSDEUNet, DBIMUNet, DDBMUNet, DDIBUNet, I2SBDiffusersUNet, LBMUNet
from src.models import FCDM, FCDM_MODELS, FCDMImageCond, FCDM_S, FCDM_B, FCDM_L, FCDM_XL
from src.models import UNSBGenerator, UNSBDiscriminator, UNSBEnergyNet
from src.models import LocalDiffusionUNet, ConditionEncoder
from src.models import Pix2PixHDGenerator, Pix2PixHDGlobalGenerator
from src.models import StarGANGenerator, StarGANDiscriminator, StarGANResidualBlock
from src.pipelines import (
    BBDMPipeline,
    BBDMPipelineOutput,
    BDBMPipeline,
    BDBMPipelineOutput,
    BiBBDMPipeline,
    BiBBDMPipelineOutput,
    CDTSDEPipeline,
    CDTSDEPipelineOutput,
    DBIMPipeline,
    DBIMPipelineOutput,
    DDBMPipeline,
    DDBMPipelineOutput,
    DDIBPipeline,
    DDIBPipelineOutput,
    ECSIPipeline,
    ECSIPipelineOutput,
    load_ecsi_pipeline,
    I2SBPipeline,
    I2SBPipelineOutput,
    LBMPipeline,
    LBMPipelineOutput,
    StegoGANPipeline,
    StegoGANPipelineOutput,
    UNSBPipeline,
    UNSBPipelineOutput,
    LocalDiffusionPipeline,
    LocalDiffusionPipelineOutput,
    Pix2PixHDPipeline,
    Pix2PixHDPipelineOutput,
    load_pix2pixhd_pipeline,
    StarGANPipeline,
    StarGANPipelineOutput,
    load_stargan_pipeline,
    ECSIPipeline,
    ECSIPipelineOutput,
    load_ecsi_pipeline,
    FCDMPipeline,
    FCDMPipelineOutput,
    FCDMImageCondPipeline,
)
from src.schedulers import (
    BBDMScheduler,
    BBDMSchedulerOutput,
    BDBMScheduler,
    BDBMSchedulerOutput,
    BiBBDMScheduler,
    BiBBDMSchedulerOutput,
    CDTSDEScheduler,
    CDTSDESchedulerOutput,
    DBIMScheduler,
    DBIMSchedulerOutput,
    DDBMScheduler,
    DDBMSchedulerOutput,
    DDIBScheduler,
    DDIBSchedulerOutput,
    I2SBScheduler,
    I2SBSchedulerOutput,
    LBMScheduler,
    LBMSchedulerOutput,
    SiD2Scheduler,
    UNSBScheduler,
    UNSBSchedulerOutput,
    LocalDiffusionScheduler,
    LocalDiffusionSchedulerOutput,
    FCDMScheduler,
    FCDMSchedulerOutput,
)

__all__ = [
    # Models — native
    "UNetGenerator",
    "ResNetGenerator",
    "PatchGANDiscriminator",
    "I2SBUNet",
    "SiTBackbone",
    "SIT_CONFIGS",
    "ResnetMaskV1Generator",
    "ResnetMaskV3Generator",
    "NetMatchability",
    # Models — UNSB
    "UNSBGenerator",
    "UNSBDiscriminator",
    "UNSBEnergyNet",
    # Models — Local Diffusion
    "LocalDiffusionUNet",
    "ConditionEncoder",
    "Pix2PixHDGenerator",
    "Pix2PixHDGlobalGenerator",
    "StarGANGenerator",
    "StarGANDiscriminator",
    "StarGANResidualBlock",
    # Models — diffusers UNet wrappers
    "BBDMUNet",
    "BDBMUNet",
    "BiBBDMUNet",
    "CDTSDEUNet",
    "DBIMUNet",
    "DDBMUNet",
    "DDIBUNet",
    "I2SBDiffusersUNet",
    "LBMUNet",
    # Models — FCDM
    "FCDM",
    "FCDM_S",
    "FCDM_B",
    "FCDM_L",
    "FCDM_XL",
    "FCDMImageCond",
    "FCDM_MODELS",
    # Schedulers
    "BBDMScheduler",
    "BBDMSchedulerOutput",
    "BDBMScheduler",
    "BDBMSchedulerOutput",
    "BiBBDMScheduler",
    "BiBBDMSchedulerOutput",
    "CDTSDEScheduler",
    "CDTSDESchedulerOutput",
    "DBIMScheduler",
    "DBIMSchedulerOutput",
    "DDBMScheduler",
    "DDBMSchedulerOutput",
    "DDIBScheduler",
    "DDIBSchedulerOutput",
    "I2SBScheduler",
    "I2SBSchedulerOutput",
    "LBMScheduler",
    "LBMSchedulerOutput",
    "SiD2Scheduler",
    "UNSBScheduler",
    "UNSBSchedulerOutput",
    "LocalDiffusionScheduler",
    "LocalDiffusionSchedulerOutput",
    "FCDMScheduler",
    "FCDMSchedulerOutput",
    # Pipelines
    "BBDMPipeline",
    "BBDMPipelineOutput",
    "BDBMPipeline",
    "BDBMPipelineOutput",
    "BiBBDMPipeline",
    "BiBBDMPipelineOutput",
    "CDTSDEPipeline",
    "CDTSDEPipelineOutput",
    "DBIMPipeline",
    "DBIMPipelineOutput",
    "DDBMPipeline",
    "DDBMPipelineOutput",
    "DDIBPipeline",
    "DDIBPipelineOutput",
    "I2SBPipeline",
    "I2SBPipelineOutput",
    "LBMPipeline",
    "LBMPipelineOutput",
    "StegoGANPipeline",
    "StegoGANPipelineOutput",
    "UNSBPipeline",
    "UNSBPipelineOutput",
    "LocalDiffusionPipeline",
    "LocalDiffusionPipelineOutput",
    "Pix2PixHDPipeline",
    "Pix2PixHDPipelineOutput",
    "load_pix2pixhd_pipeline",
    "StarGANPipeline",
    "StarGANPipelineOutput",
    "load_stargan_pipeline",
    "ECSIPipeline",
    "ECSIPipelineOutput",
    "load_ecsi_pipeline",
    "FCDMPipeline",
    "FCDMPipelineOutput",
    "FCDMImageCondPipeline",
    # Data
    "PairedImageDataset",
    "UnpairedImageDataset",
    "get_transforms",
    "default_transforms",
    # Losses
    "GANLoss",
    "PerceptualLoss",
    # Inference
    "ImageTranslator",
    # Metrics
    "PairedImageMetricEvaluator",
    "compute_psnr",
    "compute_ssim",
    "compute_lpips",
    "compute_dists",
    "compute_fid",
]
