pytorch-image-translation-models 0.1.0

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

pytorch-image-translation-models

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

Installation

Install from source

pip install -e .

With optional dependencies:

# With training extras (accelerate, peft, datasets, tensorboard)
pip install -e ".[training]"

# With metrics extras (torchmetrics, lpips, torch-fidelity, scipy)
pip install -e ".[metrics]"

# Everything
pip install -e ".[all]"

Note: PyTorch is listed as a dependency but you may want to install a specific CUDA build first. See PyTorch — Get Started for details.

Features

Models

• GAN generators — UNetGenerator (encoder-decoder with skip connections), ResNetGenerator (residual blocks)
• GAN discriminators — PatchGANDiscriminator (Markovian patch-level classifier)
• Diffusion bridge — I2SBUNet (ADM-style U-Net for Image-to-Image Schrödinger Bridge)

Schedulers

• I2SBScheduler — Symmetric beta schedule with forward/reverse bridge kernels for I2SB

Pipelines

• I2SBPipeline — End-to-end inference for I2SB models (supports "pt", "pil", "np" output)

Data

• PairedImageDataset / UnpairedImageDataset with configurable transform pipelines

Losses

• GANLoss (vanilla / LSGAN / hinge), VGG-based PerceptualLoss

Training

• Pix2PixTrainer — Paired GAN training with checkpoint save/load
• I2SBTrainer — I2SB bridge model training (in examples/i2sb/)

Metrics

• compute_psnr, compute_ssim, compute_lpips, compute_fid

Quick Start

GAN-based translation (Pix2Pix)

import src

gen = src.UNetGenerator(in_channels=3, out_channels=3)
disc = src.PatchGANDiscriminator(in_channels=6)

from src.training import Pix2PixTrainer, TrainingConfig
config = TrainingConfig(epochs=100, device="cuda")
trainer = Pix2PixTrainer(gen, disc, config)
trainer.fit(dataloader)  # expects {"source": tensor, "target": tensor}

translator = src.ImageTranslator(gen, device="cuda")
result = translator.predict(pil_image)

Diffusion bridge translation (I2SB)

from src.models.unet import I2SBUNet, create_model
from src.schedulers import I2SBScheduler
from src.pipelines.i2sb import I2SBPipeline

# Create model and scheduler
model = create_model(
    image_size=256, in_channels=3, num_channels=128,
    num_res_blocks=2, attention_resolutions="32,16,8",
    condition_mode="concat",
)
scheduler = I2SBScheduler(interval=1000, beta_max=0.3)

# Inference pipeline
pipeline = I2SBPipeline(unet=model, scheduler=scheduler)
result = pipeline(source_tensor, nfe=20, output_type="pt")

I2SB training with task configs

from examples.i2sb.config import sar2eo_config
from examples.i2sb.trainer import I2SBTrainer

cfg = sar2eo_config(resolution=256, train_batch_size=8)
trainer = I2SBTrainer(cfg)
model = trainer.build_model()
scheduler = trainer.build_scheduler()

# Single-step loss computation
loss = I2SBTrainer.compute_training_loss(model, scheduler, source_batch, target_batch)
loss.backward()

Package Structure

src/
├── __init__.py              # Public API
├── models/
│   ├── generators.py        # UNetGenerator, ResNetGenerator
│   ├── discriminators.py    # PatchGANDiscriminator
│   └── unet/                # ADM-style U-Net for I2SB
│       ├── i2sb_unet.py     # I2SBUNet
│       └── unet_2d.py       # create_model factory
├── schedulers/
│   └── i2sb.py              # I2SBScheduler
├── pipelines/
│   └── i2sb.py              # I2SBPipeline
├── data/
│   ├── datasets.py          # PairedImageDataset, UnpairedImageDataset
│   └── transforms.py        # get_transforms, default_transforms
├── losses/
│   ├── adversarial.py       # GANLoss
│   └── perceptual.py        # PerceptualLoss
├── training/
│   └── trainer.py           # Pix2PixTrainer, TrainingConfig
├── inference/
│   └── predictor.py         # ImageTranslator
└── metrics/
    └── image_quality.py     # PSNR, SSIM, LPIPS, FID
examples/
└── i2sb/
    ├── config.py            # TaskConfig, sar2eo_config, etc.
    └── trainer.py           # I2SBTrainer

Credits

Reference papers

• I2SB: Image-to-Image Schrödinger Bridge (ICML 2023)
• DDBM: Denoising Diffusion Bridge Models (ICLR 2024)
• DDIB: Dual Diffusion Implicit Bridges (ICLR 2023)
• BBDM: Image-to-Image Translation with Brownian Bridge Diffusion Models (CVPR 2023)
• CUT: Contrastive Unpaired Translation (ECCV 2020)
• CycleGAN (ICCV 2017)
• img2img-turbo (2024)

License

MIT