"""PyTorch image-to-image translation models.

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

__version__ = "0.2.9"

from src.data import PairedImageDataset, UnpairedImageDataset, default_transforms, get_transforms
from src.inference 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 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,
)
from src.schedulers import (
    BBDMScheduler,
    BBDMSchedulerOutput,
    BDBMScheduler,
    BDBMSchedulerOutput,
    BiBBDMScheduler,
    BiBBDMSchedulerOutput,
    CDTSDEScheduler,
    CDTSDESchedulerOutput,
    DBIMScheduler,
    DBIMSchedulerOutput,
    DDBMScheduler,
    DDBMSchedulerOutput,
    DDIBScheduler,
    DDIBSchedulerOutput,
    I2SBScheduler,
    I2SBSchedulerOutput,
    LBMScheduler,
    LBMSchedulerOutput,
    UNSBScheduler,
    UNSBSchedulerOutput,
    LocalDiffusionScheduler,
    LocalDiffusionSchedulerOutput,
)

__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",
    # Schedulers
    "BBDMScheduler",
    "BBDMSchedulerOutput",
    "BDBMScheduler",
    "BDBMSchedulerOutput",
    "BiBBDMScheduler",
    "BiBBDMSchedulerOutput",
    "CDTSDEScheduler",
    "CDTSDESchedulerOutput",
    "DBIMScheduler",
    "DBIMSchedulerOutput",
    "DDBMScheduler",
    "DDBMSchedulerOutput",
    "DDIBScheduler",
    "DDIBSchedulerOutput",
    "I2SBScheduler",
    "I2SBSchedulerOutput",
    "LBMScheduler",
    "LBMSchedulerOutput",
    "UNSBScheduler",
    "UNSBSchedulerOutput",
    "LocalDiffusionScheduler",
    "LocalDiffusionSchedulerOutput",
    # 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",
    # Data
    "PairedImageDataset",
    "UnpairedImageDataset",
    "get_transforms",
    "default_transforms",
    # Losses
    "GANLoss",
    "PerceptualLoss",
    # Inference
    "ImageTranslator",
    # Metrics
    "PairedImageMetricEvaluator",
    "compute_psnr",
    "compute_ssim",
    "compute_lpips",
    "compute_dists",
    "compute_fid",
]
