pytorch-image-translation-models

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

Installation

Install from PyPI

pip install pytorch-image-translation-models

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.

Quick Start

Examples default to device="cuda". If your environment is CPU-only, replace "cuda" with "cpu".

from PIL import Image

# Baseline method (UNSB)
from src.pipelines.unsb import UNSBPipeline
unsb = UNSBPipeline.from_pretrained(
    "path/to/UNSB-ckpt/horse2zebra",  # https://huggingface.co/BiliSakura/UNSB-ckpt
    subfolder="generator",
    scheduler_num_timesteps=5,
    scheduler_tau=0.01,
)
unsb.to("cuda")
unsb_out = unsb(source_image=source, output_type="pil")
unsb_out.images[0].save("unsb_output.png")

# Community method (DiffuseIT) - text/image-guided diffusion translation
from examples.community.diffuseit import load_diffuseit_community_pipeline

pipe = load_diffuseit_community_pipeline(
    "/path/to/BiliSakura/DiffuseIT-ckpt/imagenet256-uncond",
)
pipe.to("cuda")
out = pipe(
    source_image=source,
    prompt="Black Leopard",
    source="Lion",
    use_range_restart=True,
    use_noise_aug_all=True,
    output_type="pil",
)
out.images[0].save("diffuseit_output.png")

# Community method (E3Diff)
from examples.community.e3diff import E3DiffPipeline
e3diff = E3DiffPipeline.from_pretrained("path/to/E3Diff-ckpt/SEN12")
e3diff.to("cuda")
community_out = e3diff(source_image=source, num_inference_steps=50, output_type="pil")
community_out.images[0].save("e3diff_output.png")

Documentation

All information regarding per-method checkpoint folder conventions required by from_pretrained(...), as well as comprehensive package documentation, is integrated below.

Credits

This repository/package is primarily built upon 4th-MAVIC-T by the EarthBridge Team:

• Zheyuan Chen — bilisakura@zju.edu.cn
• Yuanshen Guan — guanys@mail.ustc.edu.cn

License

MIT