This repository contains a minimal implementation of Domain Guidance, adapted from the DiT repository.
- 🚩 2025.04.14: A quick implementation is released.
- 🚩 2025.04.02: ICLR 2025 camera-ready version is released on arXiv.
The main idea is to replace the unconditional branch in the CFG (Classifier-Free Guidance) implementation.
A quick overview of the key code is provided below. You can customize your own implementation following this example:
def build_cfg_forward_fn(c_model, u_model):
def cfg_forward_fn(x, t, y, cfg_scale):
c_half = x[: len(x) // 2]
c_half_t = t[: len(t) // 2]
c_half_y = y[: len(y) // 2]
u_half_x = c_half
u_half_t = t[len(t) // 2 :]
u_half_y = y[len(y) // 2 :]
c_out = c_model(c_half, c_half_t, c_half_y)
u_out = u_model(u_half_x, u_half_t, u_half_y)
c_eps, c_rest = c_out[:, :c_model.in_channels], c_out[:, c_model.in_channels:]
u_eps, u_rest = u_out[:, :u_model.in_channels], u_out[:, u_model.in_channels:]
half_eps = u_eps + cfg_scale * (c_eps - u_eps)
eps = torch.cat([half_eps, half_eps], dim=0)
rest = torch.cat([c_rest, u_rest], dim=0)
return torch.cat([eps, rest], dim=1)
return cfg_forward_fn@inproceedings{zhong2025domain,
title={Domain Guidance: A Simple Transfer Approach for a Pre-trained Diffusion Model},
author={Jincheng Zhong and XiangCheng Zhang and Jianmin Wang and Mingsheng Long},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=PplM2kDrl3}
}If you have any question, please contact [email protected].