Initial release: FUDA custom nodes for ComfyUI
Implements Fourier-based Unsupervised Domain Adaptation from Wang et al., Knowledge-Based Systems, 2023. Three nodes provided: - FUDA Image Adaptation (core amplitude mixing) - FUDA + Channel Attention (with FTCA) - FUDA Amplitude Visualiser Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
bemoregt
@@ -0,0 +1,173 @@
|
||||
# ComfyUI-FUDA
|
||||
|
||||
**Fourier-based Unsupervised Domain Adaptation nodes for ComfyUI.**
|
||||
|
||||
Implements the FUDA method from:
|
||||
|
||||
> *Boosting unsupervised domain adaptation: A Fourier approach*
|
||||
> Mengzhu Wang, Shanshan Wang, Ye Wang, Wei Wang, Tianyi Liang, Junyang Chen, Zhigang Luo
|
||||
> Knowledge-Based Systems, 2023. https://doi.org/10.1016/j.knosys.2023.110325
|
||||
|
||||
---
|
||||
|
||||
## How It Works
|
||||
|
||||
The Fourier transform decomposes an image into two components:
|
||||
|
||||
- **Amplitude** — encodes low-level statistics: colour, brightness, and style.
|
||||
- **Phase** — encodes high-level semantics: structure, edges, and content.
|
||||
|
||||
FUDA exploits this property to transfer the visual style of a **reference** image onto a **source** image while keeping its content intact:
|
||||
|
||||
1. Apply 2-D FFT to both images and shift the DC component to the centre.
|
||||
2. Mix the **low-frequency amplitude** of the reference into the source using a weighted blend controlled by `beta` (region size) and `alpha` (blend strength).
|
||||
3. Reconstruct via inverse FFT using the mixed amplitude and the original source phase.
|
||||
|
||||
An optional **Fourier Transform Channel Attention (FTCA)** module further recalibrates channel responses based on per-channel spectral energy, as described in the paper.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
All nodes live under the **`FUDA`** category in the ComfyUI node menu.
|
||||
|
||||
---
|
||||
|
||||
### FUDA Image Adaptation
|
||||
|
||||
The core node. Transfers the low-level style of `reference_image` onto `source_image`.
|
||||
|
||||
| Port | Type | Description |
|
||||
|---|---|---|
|
||||
| `source_image` | IMAGE (input) | Image to be style-adapted |
|
||||
| `reference_image` | IMAGE (input) | Target-domain reference image |
|
||||
| `beta` | FLOAT 0.001–0.5 | Low-frequency band radius as a fraction of the shorter spatial dimension. Small values (0.01–0.1) affect global colour and tone; larger values also mix mid-frequency texture. **Default: 0.09** |
|
||||
| `alpha` | FLOAT 0–1 | Blend weight for the reference amplitude. `0` = no change, `1` = full reference style. **Default: 0.5** |
|
||||
| `adapted_image` | IMAGE (output) | Style-adapted result |
|
||||
|
||||
**Tips:**
|
||||
- Start with `beta = 0.09`, `alpha = 0.5` and adjust from there.
|
||||
- Increasing `beta` broadens the spectral region mixed — more texture change but potentially more artefacts.
|
||||
- `alpha = 1.0` reproduces the original FDA (Fourier Domain Adaptation) method.
|
||||
|
||||
---
|
||||
|
||||
### FUDA + Channel Attention
|
||||
|
||||
Extends the core node with a **Fourier Transform Channel Attention (FTCA)** pass. After amplitude mixing, FTCA computes per-channel spectral energy weights through a small MLP and recalibrates the adapted image to sharpen discriminative features.
|
||||
|
||||
| Port | Type | Description |
|
||||
|---|---|---|
|
||||
| `source_image` | IMAGE (input) | Image to be style-adapted |
|
||||
| `reference_image` | IMAGE (input) | Target-domain reference image |
|
||||
| `beta` | FLOAT 0.001–0.5 | Low-frequency band ratio (see above) |
|
||||
| `alpha` | FLOAT 0–1 | Reference amplitude blend weight |
|
||||
| `attention_strength` | FLOAT 0–1 | Blend between raw adaptation (`0`) and FTCA-refined output (`1`). **Default: 0.5** |
|
||||
| `reduction` | INT 1–16 | Channel reduction ratio inside the attention MLP. Higher = fewer parameters. **Default: 4** |
|
||||
| `adapted_image` | IMAGE (output) | Attention-refined, style-adapted result |
|
||||
|
||||
---
|
||||
|
||||
### FUDA Amplitude Visualiser
|
||||
|
||||
Renders the centred Fourier amplitude spectrum of an image as a greyscale heatmap. Useful for comparing domain differences before and after adaptation.
|
||||
|
||||
| Port | Type | Description |
|
||||
|---|---|---|
|
||||
| `image` | IMAGE (input) | Any image |
|
||||
| `log_scale` | BOOLEAN | Apply `log(1 + amplitude)` compression for better dynamic range visualisation. **Default: true** |
|
||||
| `amplitude_map` | IMAGE (output) | Normalised amplitude heatmap (bright = high energy) |
|
||||
|
||||
---
|
||||
|
||||
## Installation
|
||||
|
||||
### Option A — Clone into `custom_nodes`
|
||||
|
||||
```bash
|
||||
cd /path/to/ComfyUI/custom_nodes
|
||||
git clone https://github.com/your-repo/ComfyUI_FUDA
|
||||
```
|
||||
|
||||
### Option B — Symlink an existing clone
|
||||
|
||||
```bash
|
||||
ln -s /Users/m1_4k/ComfyUI_FUDA /path/to/ComfyUI/custom_nodes/ComfyUI_FUDA
|
||||
```
|
||||
|
||||
Restart ComfyUI. The three nodes will appear under the **FUDA** category.
|
||||
|
||||
### Requirements
|
||||
|
||||
No additional packages are needed beyond what ComfyUI already installs.
|
||||
|
||||
| Dependency | Version |
|
||||
|---|---|
|
||||
| Python | ≥ 3.10 |
|
||||
| PyTorch | ≥ 2.0 (for `torch.fft`) |
|
||||
| NumPy | any recent |
|
||||
|
||||
---
|
||||
|
||||
## Example Workflow
|
||||
|
||||
```
|
||||
[Load Image] ──► source_image ─┐
|
||||
├─► [FUDA Image Adaptation] ──► [Preview Image]
|
||||
[Load Image] ──► reference_image ─┘ beta=0.09 alpha=0.5
|
||||
```
|
||||
|
||||
To inspect the effect in the frequency domain:
|
||||
|
||||
```
|
||||
[Load Image] ──► [FUDA Amplitude Visualiser] ──► [Preview Image] (before)
|
||||
[adapted_image] ──► [FUDA Amplitude Visualiser] ──► [Preview Image] (after)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Parameter Guide
|
||||
|
||||
| Goal | Suggested settings |
|
||||
|---|---|
|
||||
| Subtle colour grading | `beta=0.03`, `alpha=0.3` |
|
||||
| Moderate style transfer | `beta=0.09`, `alpha=0.5` |
|
||||
| Strong style transfer | `beta=0.2`, `alpha=0.8` |
|
||||
| Maximum reference style | `beta=0.3`, `alpha=1.0` |
|
||||
| Add attention refinement | Use **FUDA + Channel Attention**, `attention_strength=0.5` |
|
||||
|
||||
---
|
||||
|
||||
## Background: Why Fourier?
|
||||
|
||||
The Fourier transform separates an image's *what* (phase → semantics) from its *how it looks* (amplitude → style). By blending only the low-frequency amplitude:
|
||||
|
||||
- The **content and structure** of the source image are fully preserved (phase is untouched).
|
||||
- The **colour palette, illumination, and global style** shift toward the reference domain.
|
||||
|
||||
This is the foundation of **FDA** (Fourier Domain Adaptation, Yang et al., 2020), which FUDA extends by:
|
||||
1. Using a **weighted blend** instead of a hard replacement.
|
||||
2. Adding **Fourier Transform Channel Attention** to capture richer spectral feature diversity.
|
||||
|
||||
---
|
||||
|
||||
## Citation
|
||||
|
||||
If you use this node in your work, please cite the original paper:
|
||||
|
||||
```bibtex
|
||||
@article{wang2023fuda,
|
||||
title = {Boosting unsupervised domain adaptation: A Fourier approach},
|
||||
author = {Wang, Mengzhu and Wang, Shanshan and Wang, Ye and Wang, Wei
|
||||
and Liang, Tianyi and Chen, Junyang and Luo, Zhigang},
|
||||
journal = {Knowledge-Based Systems},
|
||||
year = {2023},
|
||||
doi = {10.1016/j.knosys.2023.110325}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,296 @@
|
||||
"""
|
||||
FUDA: Fourier-based Unsupervised Domain Adaptation for ComfyUI
|
||||
Based on: "Boosting unsupervised domain adaptation: A Fourier approach"
|
||||
Wang et al., Knowledge-Based Systems, 2023
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core FUDA transforms
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def fuda_amplitude_mix(src: torch.Tensor, ref: torch.Tensor,
|
||||
beta: float = 0.1, alpha: float = 0.5) -> torch.Tensor:
|
||||
"""
|
||||
Fourier Domain Adaptation: fuse low-frequency amplitude of reference
|
||||
into source, keeping source phase intact.
|
||||
|
||||
Args:
|
||||
src : [B, C, H, W] float tensor in [0, 1]
|
||||
ref : [B, C, H, W] float tensor in [0, 1] (same spatial size as src)
|
||||
beta : fraction of the spectrum (from DC) to mix [0, 0.5]
|
||||
alpha: blend weight for reference amplitude [0, 1]
|
||||
Returns:
|
||||
adapted [B, C, H, W] float tensor
|
||||
"""
|
||||
B, C, H, W = src.shape
|
||||
|
||||
fft_src = torch.fft.fft2(src)
|
||||
fft_ref = torch.fft.fft2(ref)
|
||||
|
||||
# shift DC to centre
|
||||
fft_src = torch.fft.fftshift(fft_src, dim=(-2, -1))
|
||||
fft_ref = torch.fft.fftshift(fft_ref, dim=(-2, -1))
|
||||
|
||||
amp_src = torch.abs(fft_src)
|
||||
pha_src = torch.angle(fft_src)
|
||||
amp_ref = torch.abs(fft_ref)
|
||||
|
||||
b = max(1, int(np.floor(min(H, W) * beta)))
|
||||
hc, wc = H // 2, W // 2
|
||||
h0, h1 = max(0, hc - b), min(H, hc + b)
|
||||
w0, w1 = max(0, wc - b), min(W, wc + b)
|
||||
|
||||
amp_new = amp_src.clone()
|
||||
amp_new[:, :, h0:h1, w0:w1] = (
|
||||
(1.0 - alpha) * amp_src[:, :, h0:h1, w0:w1]
|
||||
+ alpha * amp_ref[:, :, h0:h1, w0:w1]
|
||||
)
|
||||
|
||||
fft_new = torch.polar(amp_new, pha_src)
|
||||
fft_new = torch.fft.ifftshift(fft_new, dim=(-2, -1))
|
||||
adapted = torch.fft.ifft2(fft_new).real
|
||||
return adapted
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fourier Transform Channel Attention (FTCA)
|
||||
# Lightweight pixel-space version: reweights channels via their spectral energy
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FourierChannelAttention(nn.Module):
|
||||
"""
|
||||
Pixel-space Fourier Channel Attention.
|
||||
Computes per-channel spectral energy weights and recalibrates the feature map.
|
||||
"""
|
||||
def __init__(self, channels: int, reduction: int = 4):
|
||||
super().__init__()
|
||||
mid = max(1, channels // reduction)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(channels, mid, bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(mid, channels, bias=False),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x: [B, C, H, W]
|
||||
fft_x = torch.fft.fft2(x)
|
||||
energy = (torch.abs(fft_x) ** 2).mean(dim=(-2, -1)) # [B, C]
|
||||
w = self.fc(energy).unsqueeze(-1).unsqueeze(-1) # [B, C, 1, 1]
|
||||
return x * w
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ComfyUI Nodes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FUDANode:
|
||||
"""
|
||||
FUDA Image Adaptation
|
||||
─────────────────────
|
||||
Adapts *source_image* toward the visual style of *reference_image* using
|
||||
Fourier-domain amplitude mixing (the core FUDA transform from the paper).
|
||||
|
||||
• beta – controls the size of the low-frequency region that is mixed.
|
||||
Small values (0.01–0.1) transfer colour/style; larger values
|
||||
can also affect mid-frequency texture.
|
||||
• alpha – blend weight: 0 = no change, 1 = full reference amplitude.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"source_image": ("IMAGE",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"beta": ("FLOAT", {
|
||||
"default": 0.09,
|
||||
"min": 0.001,
|
||||
"max": 0.5,
|
||||
"step": 0.001,
|
||||
"display": "slider",
|
||||
"tooltip": "Low-frequency band ratio (0.01–0.5). Higher → wider spectral mix.",
|
||||
}),
|
||||
"alpha": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"display": "slider",
|
||||
"tooltip": "Reference amplitude weight. 0=no change, 1=full reference style.",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("adapted_image",)
|
||||
FUNCTION = "adapt"
|
||||
CATEGORY = "FUDA"
|
||||
DESCRIPTION = "Fourier Domain Adaptation – adapts source image style toward a reference domain."
|
||||
|
||||
def adapt(self, source_image: torch.Tensor, reference_image: torch.Tensor,
|
||||
beta: float, alpha: float):
|
||||
# ComfyUI tensors: [B, H, W, C] float32
|
||||
src = source_image.permute(0, 3, 1, 2) # → [B, C, H, W]
|
||||
ref = reference_image.permute(0, 3, 1, 2)
|
||||
|
||||
# broadcast: match batch dims
|
||||
if src.shape[0] != ref.shape[0]:
|
||||
if ref.shape[0] == 1:
|
||||
ref = ref.expand(src.shape[0], -1, -1, -1)
|
||||
elif src.shape[0] == 1:
|
||||
src = src.expand(ref.shape[0], -1, -1, -1)
|
||||
|
||||
# Resize reference spatially if needed
|
||||
if src.shape[-2:] != ref.shape[-2:]:
|
||||
ref = F.interpolate(ref, size=src.shape[-2:],
|
||||
mode="bilinear", align_corners=False)
|
||||
|
||||
adapted = fuda_amplitude_mix(src, ref, beta=beta, alpha=alpha)
|
||||
adapted = adapted.clamp(0.0, 1.0).permute(0, 2, 3, 1) # → [B, H, W, C]
|
||||
return (adapted,)
|
||||
|
||||
|
||||
class FUDAWithAttentionNode:
|
||||
"""
|
||||
FUDA Image Adaptation + Fourier Channel Attention
|
||||
──────────────────────────────────────────────────
|
||||
Same as FUDANode but also applies a Fourier-based channel attention pass
|
||||
to enhance discriminative features in the adapted result.
|
||||
|
||||
Extra parameters:
|
||||
• attention_strength – blends the attention-refined output with the raw
|
||||
adapted output (0 = skip attention, 1 = full FTCA).
|
||||
• reduction – channel reduction ratio in the attention MLP.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"source_image": ("IMAGE",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"beta": ("FLOAT", {
|
||||
"default": 0.09,
|
||||
"min": 0.001, "max": 0.5, "step": 0.001,
|
||||
"display": "slider",
|
||||
}),
|
||||
"alpha": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"display": "slider",
|
||||
}),
|
||||
"attention_strength": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"display": "slider",
|
||||
"tooltip": "0 = no attention, 1 = full FTCA recalibration.",
|
||||
}),
|
||||
"reduction": ("INT", {
|
||||
"default": 4,
|
||||
"min": 1, "max": 16, "step": 1,
|
||||
"tooltip": "Channel reduction ratio for attention MLP.",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("adapted_image",)
|
||||
FUNCTION = "adapt_with_attention"
|
||||
CATEGORY = "FUDA"
|
||||
DESCRIPTION = "FUDA + Fourier Channel Attention for enhanced style adaptation."
|
||||
|
||||
def adapt_with_attention(self, source_image: torch.Tensor,
|
||||
reference_image: torch.Tensor,
|
||||
beta: float, alpha: float,
|
||||
attention_strength: float, reduction: int):
|
||||
src = source_image.permute(0, 3, 1, 2)
|
||||
ref = reference_image.permute(0, 3, 1, 2)
|
||||
|
||||
if src.shape[0] != ref.shape[0]:
|
||||
if ref.shape[0] == 1:
|
||||
ref = ref.expand(src.shape[0], -1, -1, -1)
|
||||
elif src.shape[0] == 1:
|
||||
src = src.expand(ref.shape[0], -1, -1, -1)
|
||||
|
||||
if src.shape[-2:] != ref.shape[-2:]:
|
||||
ref = F.interpolate(ref, size=src.shape[-2:],
|
||||
mode="bilinear", align_corners=False)
|
||||
|
||||
# Step 1: Fourier amplitude adaptation
|
||||
adapted = fuda_amplitude_mix(src, ref, beta=beta, alpha=alpha)
|
||||
|
||||
# Step 2: Fourier Channel Attention
|
||||
C = adapted.shape[1]
|
||||
ftca = FourierChannelAttention(channels=C, reduction=reduction)
|
||||
ftca.eval()
|
||||
with torch.no_grad():
|
||||
refined = ftca(adapted)
|
||||
|
||||
# Blend
|
||||
out = (1.0 - attention_strength) * adapted + attention_strength * refined
|
||||
out = out.clamp(0.0, 1.0).permute(0, 2, 3, 1)
|
||||
return (out,)
|
||||
|
||||
|
||||
class FUDAAmplitudeVisNode:
|
||||
"""
|
||||
FUDA Amplitude Visualiser
|
||||
─────────────────────────
|
||||
Outputs a visual representation of the Fourier amplitude spectrum of an
|
||||
image. Useful for understanding domain differences.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"log_scale": ("BOOLEAN", {"default": True,
|
||||
"tooltip": "Apply log(1+amp) for better visualisation."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("amplitude_map",)
|
||||
FUNCTION = "visualise"
|
||||
CATEGORY = "FUDA"
|
||||
DESCRIPTION = "Visualise Fourier amplitude spectrum (shifted, per-channel mean)."
|
||||
|
||||
def visualise(self, image: torch.Tensor, log_scale: bool):
|
||||
x = image.permute(0, 3, 1, 2) # [B,C,H,W]
|
||||
fft_x = torch.fft.fftshift(torch.fft.fft2(x), dim=(-2, -1))
|
||||
amp = torch.abs(fft_x).mean(dim=1, keepdim=True) # [B,1,H,W]
|
||||
if log_scale:
|
||||
amp = torch.log1p(amp)
|
||||
# Normalise per image to [0,1]
|
||||
B = amp.shape[0]
|
||||
amp_flat = amp.view(B, -1)
|
||||
mn = amp_flat.min(dim=1).values.view(B, 1, 1, 1)
|
||||
mx = amp_flat.max(dim=1).values.view(B, 1, 1, 1)
|
||||
amp = (amp - mn) / (mx - mn + 1e-8)
|
||||
amp = amp.expand(-1, 3, -1, -1).permute(0, 2, 3, 1) # [B,H,W,3]
|
||||
return (amp.clamp(0, 1),)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Node registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FUDANode": FUDANode,
|
||||
"FUDAWithAttentionNode": FUDAWithAttentionNode,
|
||||
"FUDAAmplitudeVisNode": FUDAAmplitudeVisNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FUDANode": "FUDA Image Adaptation",
|
||||
"FUDAWithAttentionNode": "FUDA + Channel Attention",
|
||||
"FUDAAmplitudeVisNode": "FUDA Amplitude Visualiser",
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
[project]
|
||||
name = "comfyui-fuda"
|
||||
version = "1.0.0"
|
||||
description = "FUDA: Fourier-based Unsupervised Domain Adaptation nodes for ComfyUI"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = []
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/your-repo/comfyui-fuda"
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = ""
|
||||
DisplayName = "FUDA – Fourier Domain Adaptation"
|
||||
Icon = ""
|
||||
Reference in New Issue
Block a user