CDC-FM (Carré du Champ Flow Matching) Implementation

.github/workflows/tests.yml

@@ -43,7 +43,7 @@ jobs:
       - name: Install dependencies
         run: |
           # Pre-install torch to pin version (requirements.txt has dependencies like transformers which requires pytorch)
-          pip install dadaptation==3.2 torch==${{ matrix.pytorch-version }} torchvision pytest==8.3.4
+          pip install dadaptation==3.2 torch==${{ matrix.pytorch-version }} torchvision pytest==8.3.4 faiss-cpu==1.12.0
           pip install -r requirements.txt
 
       - name: Test with pytest

.gitignore

@@ -11,3 +11,4 @@ GEMINI.md
 .claude
 .gemini
 MagicMock
+benchmark_*.py

flux_train_network.py

@@ -1,7 +1,5 @@
 import argparse
 import copy
-import math
-import random
 from typing import Any, Optional, Union
 
 import torch
@@ -36,6 +34,7 @@ def __init__(self):
         self.is_schnell: Optional[bool] = None
         self.is_swapping_blocks: bool = False
         self.model_type: Optional[str] = None
+        self.gamma_b_dataset = None  # CDC-FM Γ_b dataset
 
     def assert_extra_args(
         self,
@@ -327,9 +326,15 @@ def get_noise_pred_and_target(
         noise = torch.randn_like(latents)
         bsz = latents.shape[0]
 
-        # get noisy model input and timesteps
+        # Get CDC parameters if enabled
+        gamma_b_dataset = self.gamma_b_dataset if (self.gamma_b_dataset is not None and "latents_npz" in batch) else None
+        latents_npz_paths = batch.get("latents_npz") if gamma_b_dataset is not None else None
+
+        # Get noisy model input and timesteps
+        # If CDC is enabled, this will transform the noise with geometry-aware covariance
         noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timesteps(
-            args, noise_scheduler, latents, noise, accelerator.device, weight_dtype
+            args, noise_scheduler, latents, noise, accelerator.device, weight_dtype,
+            gamma_b_dataset=gamma_b_dataset, latents_npz_paths=latents_npz_paths
         )
 
         # pack latents and get img_ids
@@ -456,6 +461,15 @@ def update_metadata(self, metadata, args):
         metadata["ss_model_prediction_type"] = args.model_prediction_type
         metadata["ss_discrete_flow_shift"] = args.discrete_flow_shift
 
+        # CDC-FM metadata
+        metadata["ss_use_cdc_fm"] = getattr(args, "use_cdc_fm", False)
+        metadata["ss_cdc_k_neighbors"] = getattr(args, "cdc_k_neighbors", None)
+        metadata["ss_cdc_k_bandwidth"] = getattr(args, "cdc_k_bandwidth", None)
+        metadata["ss_cdc_d_cdc"] = getattr(args, "cdc_d_cdc", None)
+        metadata["ss_cdc_gamma"] = getattr(args, "cdc_gamma", None)
+        metadata["ss_cdc_adaptive_k"] = getattr(args, "cdc_adaptive_k", None)
+        metadata["ss_cdc_min_bucket_size"] = getattr(args, "cdc_min_bucket_size", None)
+
     def is_text_encoder_not_needed_for_training(self, args):
         return args.cache_text_encoder_outputs and not self.is_train_text_encoder(args)
 
@@ -494,7 +508,7 @@ def forward(hidden_states):
                         module.forward = forward_hook(module)
 
             if flux_utils.get_t5xxl_actual_dtype(text_encoder) == torch.float8_e4m3fn and text_encoder.dtype == weight_dtype:
-                logger.info(f"T5XXL already prepared for fp8")
+                logger.info("T5XXL already prepared for fp8")
             else:
                 logger.info(f"prepare T5XXL for fp8: set to {te_weight_dtype}, set embeddings to {weight_dtype}, add hooks")
                 text_encoder.to(te_weight_dtype)  # fp8
@@ -533,6 +547,72 @@ def setup_parser() -> argparse.ArgumentParser:
         help="[Deprecated] This option is deprecated. Please use `--blocks_to_swap` instead."
         " / このオプションは非推奨です。代わりに`--blocks_to_swap`を使用してください。",
     )
+
+    # CDC-FM arguments
+    parser.add_argument(
+        "--use_cdc_fm",
+        action="store_true",
+        help="Enable CDC-FM (Carré du Champ Flow Matching) for geometry-aware noise during training"
+        " / CDC-FM（Carré du Champ Flow Matching）を有効にして幾何学的ノイズを使用",
+    )
+    parser.add_argument(
+        "--cdc_k_neighbors",
+        type=int,
+        default=256,
+        help="Number of neighbors for k-NN graph in CDC-FM (default: 256)"
+        " / CDC-FMのk-NNグラフの近傍数（デフォルト: 256）",
+    )
+    parser.add_argument(
+        "--cdc_k_bandwidth",
+        type=int,
+        default=8,
+        help="Number of neighbors for bandwidth estimation in CDC-FM (default: 8)"
+        " / CDC-FMの帯域幅推定の近傍数（デフォルト: 8）",
+    )
+    parser.add_argument(
+        "--cdc_d_cdc",
+        type=int,
+        default=8,
+        help="Dimension of CDC subspace (default: 8)"
+        " / CDCサブ空間の次元（デフォルト: 8）",
+    )
+    parser.add_argument(
+        "--cdc_gamma",
+        type=float,
+        default=1.0,
+        help="CDC strength parameter (default: 1.0)"
+        " / CDC強度パラメータ（デフォルト: 1.0）",
+    )
+    parser.add_argument(
+        "--force_recache_cdc",
+        action="store_true",
+        help="Force recompute CDC cache even if valid cache exists"
+        " / 有効なCDCキャッシュが存在してもCDCキャッシュを再計算",
+    )
+    parser.add_argument(
+        "--cdc_debug",
+        action="store_true",
+        help="Enable verbose CDC debug output showing bucket details"
+        " / CDCの詳細デバッグ出力を有効化（バケット詳細表示）",
+    )
+    parser.add_argument(
+        "--cdc_adaptive_k",
+        action="store_true",
+        help="Use adaptive k_neighbors based on bucket size. If enabled, buckets smaller than k_neighbors will use "
+        "k=bucket_size-1 instead of skipping CDC entirely. Buckets smaller than cdc_min_bucket_size are still skipped."
+        " / バケットサイズに基づいてk_neighborsを適応的に調整。有効にすると、k_neighbors未満のバケットは"
+        "CDCをスキップせずk=バケットサイズ-1を使用。cdc_min_bucket_size未満のバケットは引き続きスキップ。",
+    )
+    parser.add_argument(
+        "--cdc_min_bucket_size",
+        type=int,
+        default=16,
+        help="Minimum bucket size for CDC computation. Buckets with fewer samples will use standard Gaussian noise. "
+        "Only relevant when --cdc_adaptive_k is enabled (default: 16)"
+        " / CDC計算の最小バケットサイズ。これより少ないサンプルのバケットは標準ガウスノイズを使用。"
+        "--cdc_adaptive_k有効時のみ関連（デフォルト: 16）",
+    )
+
     return parser