Upload folder using huggingface_hub

README.md

@@ -0,0 +1,158 @@
+---
+license: apache-2.0
+tags:
+  - defect-generation
+  - anomaly-detection
+  - industrial-inspection
+  - lora
+  - flux
+  - diffusion
+  - rlhf
+language: en
+pipeline_tag: image-to-image
+---
+
+# UniDG-RFT-LoRA
+
+LoRA weights for **UniDG** (Universal Defect Generation), trained via **Consistency-RFT** with Flow-GRPO and dual reward models on the UDG dataset (300K quadruplets).
+
+[[Paper]](https://arxiv.org/abs/2604.08915) [[Code]](https://github.com/RetoFan233/UniDG) [[UniDG-SFT-LoRA]](https://huggingface.co/retofan23333/UniDG-SFT-LoRA-Release)
+
+## Overview
+
+UniDG is a universal defect generation foundation model that transfers defects from a reference image to a target region via **Defect-Context Editing** and **MM-DiT multimodal attention**, without per-category fine-tuning. This checkpoint is the **Consistency-RFT** variant, further refined from UniDG-SFT using Flow-GRPO with dual reward models (Defect-Und-Reward & Defect-Recog-Reward) for improved defect fidelity and consistency.
+
+| Variant | Training | Focus |
+|---------|----------|-------|
+| UniDG-SFT | Diversity-SFT with complementary sampling | Diverse defect patterns |
+| **UniDG-RFT** (this) | Consistency-RFT with Flow-GRPO + dual rewards | Consistent & faithful defects |
+
+## Important: Usage Difference from UniDG-SFT-LoRA
+
+**The UniDG-RFT-LoRA weights are stored in PEFT format** (`adapter_model.safetensors` + `adapter_config.json`), which is different from UniDG-SFT-LoRA (which uses `pytorch_lora_weights.safetensors`). This means:
+
+- **UniDG-SFT-LoRA** can be directly loaded via the `lora_weights_path` parameter in `ImageUniDG`.
+- **UniDG-RFT-LoRA** must first be **merged into the base SFT model** using the provided `combine_peft_weights.py` script. After merging, the resulting model can be loaded directly without any additional LoRA loading step.
+
+## Repository Contents
+
+| File | Description |
+|------|-------------|
+| `adapter_model.safetensors` | PEFT LoRA weights (Consistency-RFT) |
+| `adapter_config.json` | LoRA configuration (rank=64, alpha=128) |
+| `combine_peft_weights.py` | Script to merge RFT LoRA into the base SFT model |
+
+## Step-by-Step Usage
+
+### Prerequisites
+
+- [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) (inpainting backbone)
+- [FLUX.1-Redux-dev](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev) (reference conditioning)
+- [UniDG-SFT-LoRA](https://huggingface.co/retofan23333/UniDG-SFT-LoRA-Release) (base SFT model — the RFT LoRA is fine-tuned on top of this)
+- [UniDG code](https://github.com/RetoFan233/UniDG) (inference framework)
+- Python dependencies: `diffusers`, `peft`, `torch`
+
+### Step 1: Prepare the Base SFT Model
+
+First, you need a base FLUX.1-Fill-dev model with UniDG-SFT-LoRA weights already merged in. If you haven't done this, you can prepare it by loading the SFT model and saving the merged weights:
+
+```python
+from diffusers import FluxFillPipeline
+import torch
+
+# Load base FLUX.1-Fill-dev
+pipe = FluxFillPipeline.from_pretrained(
+    "path/to/FLUX.1-Fill-dev",
+    torch_dtype=torch.bfloat16,
+)
+
+# Load SFT LoRA weights
+pipe.load_lora_weights("path/to/UniDG-SFT-LoRA-Release/pytorch_lora_weights.safetensors")
+
+# Save the merged SFT model as the base for RFT merging
+pipe.save_pretrained("path/to/FLUX.1-Fill-dev-UDG-SFT", safe_serialization=True, max_shard_size="5GB")
+```
+
+### Step 2: Merge RFT LoRA into the Base SFT Model
+
+Use the provided `combine_peft_weights.py` to merge the RFT LoRA weights into the base SFT model:
+
+```bash
+python combine_peft_weights.py \
+    --base_model_path path/to/FLUX.1-Fill-dev-UDG-SFT \
+    --lora_weights_path path/to/UniDG-RFT-LoRA-Release \
+    --output_path path/to/FLUX.1-Fill-dev-UDG-RFT \
+    --save_full_pipeline
+```
+
+Parameters:
+- `--base_model_path`: Path to the base SFT model (from Step 1)
+- `--lora_weights_path`: Path to this RFT LoRA repository (containing `adapter_model.safetensors` and `adapter_config.json`)
+- `--output_path`: Output path for the merged model
+- `--save_full_pipeline`: Save the full pipeline (including VAE, text encoder, etc.) so you can load it directly later
+- `--dtype`: Data type, default `bfloat16`
+- `--device`: Device for loading, default `cpu` (recommended to avoid OOM)
+
+> **Tip**: Use `--device cpu` (default) to save GPU memory during the merge process. The merge only needs to run once.
+
+### Step 3: Use the Merged Model with UniDG
+
+After merging, the model can be used directly with the UniDG inference code — **no additional LoRA loading is needed**:
+
+```python
+from unidg import ImageUniDG
+from PIL import Image
+import torch
+
+# Load the merged RFT model — set lora_weights_path="" since LoRA is already merged
+model = ImageUniDG(
+    flux_model_path="path/to/FLUX.1-Fill-dev-UDG-RFT",
+    redux_model_path="path/to/FLUX.1-Redux-dev",
+    lora_weights_path="",  # No additional LoRA needed!
+    device="cuda:0",
+    dtype=torch.bfloat16,
+)
+
+result, mask = model.process_images(
+    target_image=Image.open("target.jpg"),
+    reference_image=Image.open("reference.jpg"),
+    reference_mask=Image.open("reference_mask.png"),
+    target_mask=Image.open("target_mask.png"),
+    num_inference_steps=28,
+    guidance_scale=3.5,
+    seed=42,
+)
+result.save("result.png")
+```
+
+### Quick Reference: SFT vs RFT Usage
+
+| | UniDG-SFT | UniDG-RFT |
+|---|-----------|-----------|
+| Weight format | `pytorch_lora_weights.safetensors` | `adapter_model.safetensors` + `adapter_config.json` |
+| Merge required? | No | Yes (with SFT base model) |
+| `lora_weights_path` | Path to SFT weights | `""` (empty, after merge) |
+| `flux_model_path` | `path/to/FLUX.1-Fill-dev` | `path/to/merged-RFT-model` |
+| Load time | LoRA loaded on-the-fly | Pre-merged, no LoRA overhead |
+
+## LoRA Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| PEFT type | LORA |
+| Rank (r) | 64 |
+| Alpha | 128 |
+| Dropout | 0.0 |
+| Target modules | `ff.net.0.proj`, `ff.net.2`, `ff_context.net.0.proj`, `proj_mlp`, `attn.to_q`, `attn.to_v`, `attn.to_add_out`, `attn.add_k_proj`, `attn.add_v_proj`, `ff_context.net.2`, `attn.add_q_proj`, `attn.to_out.0`, `attn.to_k` |
+| Base model | FLUX.1-Fill-dev + UniDG-SFT-LoRA |
+
+## Citation
+
+```bibtex
+@article{fan2026unidg,
+  title={Large-Scale Universal Defect Generation: Foundation Models and Datasets},
+  author={Fan, Yuanting and Liu, Jun and Gao, Bin-Bin and Chen, Xiaochen and Lin, Yuhuan and Dai, Zhewei and Zhan, Jiawei and Wang, Chengjie},
+  journal={arXiv preprint arXiv:2604.08915},
+  year={2026}
+}
+```

adapter_config.json

@@ -0,0 +1,51 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "FluxTransformer2DModel",
+    "parent_library": "diffusers.models.transformers.transformer_flux"
+  },
+  "base_model_name_or_path": null,
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 128,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "qalora_group_size": 16,
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "ff.net.0.proj",
+    "ff.net.2",
+    "ff_context.net.0.proj",
+    "proj_mlp",
+    "attn.to_q",
+    "attn.to_v",
+    "attn.to_add_out",
+    "attn.add_k_proj",
+    "attn.add_v_proj",
+    "ff_context.net.2",
+    "attn.add_q_proj",
+    "attn.to_out.0",
+    "attn.to_k"
+  ],
+  "target_parameters": null,
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffdb61c5049d61ff469601d1f4a84a983d8b0aba0c2aa3c05b1da75bcfc887e7
+size 433431520

combine_peft_weights.py

@@ -0,0 +1,349 @@
+"""
+将 PEFT LoRA 权重合并到基础 Flux Transformer 模型中
+
+功能：
+1. 加载基础 Flux Fill 模型的 Transformer
+2. 加载 RL 训练的 PEFT LoRA 权重
+3. 将 LoRA 权重合并到基础模型
+4. 保存合并后的完整模型
+
+使用方法：
+python combine_peft_weights.py \
+    --base_model_path /path/to/base/model \
+    --lora_weights_path /path/to/lora/weights \
+    --output_path /path/to/output \
+    --save_full_pipeline  # 可选：保存完整 pipeline 而不只是 transformer
+"""
+
+import torch
+import argparse
+import os
+from pathlib import Path
+from diffusers import FluxFillPipeline
+from peft import PeftModel
+
+
+def merge_and_save_transformer(
+    base_model_path: str,
+    lora_weights_path: str,
+    output_path: str,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str = "cpu"
+):
+    """
+    合并 LoRA 权重到 Transformer 并保存
+    
+    Args:
+        base_model_path: 基础 Flux Fill 模型路径
+        lora_weights_path: PEFT LoRA 权重路径
+        output_path: 输出路径（保存合并后的 transformer）
+        dtype: 数据类型
+        device: 加载设备（建议用 CPU 以节省显存）
+    """
+    print("=" * 80)
+    print("Step 1: 加载基础 Flux Fill 模型...")
+    print("=" * 80)
+    
+    # 加载基础模型（只加载 transformer 部分以节省内存）
+    pipe = FluxFillPipeline.from_pretrained(
+        base_model_path,
+        torch_dtype=dtype,
+        low_cpu_mem_usage=True
+    )
+    
+    print(f"✓ 基础模型加载完成: {base_model_path}")
+    print(f"  Transformer 参数量: {sum(p.numel() for p in pipe.transformer.parameters()) / 1e9:.2f}B")
+    
+    # 移动到指定设备
+    if device != "cpu":
+        print(f"  移动 transformer 到 {device}...")
+        pipe.transformer = pipe.transformer.to(device)
+    
+    print("\n" + "=" * 80)
+    print("Step 2: 加载 PEFT LoRA 权重...")
+    print("=" * 80)
+    
+    # 加载 PEFT 模型
+    print(f"  从 {lora_weights_path} 加载 LoRA 权重...")
+    peft_model = PeftModel.from_pretrained(
+        pipe.transformer,
+        lora_weights_path,
+        is_trainable=False
+    )
+    peft_model.set_adapter("default")
+    
+    print(f"✓ PEFT 模型加载完成")
+    
+    # 检查 LoRA 配置
+    lora_config = peft_model.peft_config.get("default", None)
+    if lora_config:
+        print(f"  LoRA 配置:")
+        print(f"    - Rank (r): {lora_config.r}")
+        print(f"    - Alpha: {lora_config.lora_alpha}")
+        print(f"    - Dropout: {lora_config.lora_dropout}")
+        print(f"    - Target modules: {lora_config.target_modules}")
+    
+    print("\n" + "=" * 80)
+    print("Step 3: 合并 LoRA 权重到基础模型...")
+    print("=" * 80)
+    
+    # 合并权重
+    merged_model = peft_model.merge_and_unload()
+    
+    print(f"✓ 权重合并完成")
+    print(f"  合并后模型参数量: {sum(p.numel() for p in merged_model.parameters()) / 1e9:.2f}B")
+    
+    print("\n" + "=" * 80)
+    print("Step 4: 保存合并后的模型...")
+    print("=" * 80)
+    
+    # 创建输出目录
+    os.makedirs(output_path, exist_ok=True)
+    
+    # 保存合并后的 transformer
+    print(f"  保存到 {output_path}...")
+    merged_model.save_pretrained(
+        output_path,
+        safe_serialization=True,  # 使用 safetensors 格式
+        max_shard_size="5GB"  # 分片大小
+    )
+    
+    print(f"✓ 模型保存完成: {output_path}")
+    
+    # 保存模型配置信息
+    info_path = os.path.join(output_path, "merge_info.txt")
+    with open(info_path, "w") as f:
+        f.write(f"Base model: {base_model_path}\n")
+        f.write(f"LoRA weights: {lora_weights_path}\n")
+        f.write(f"Merged model: {output_path}\n")
+        f.write(f"Data type: {dtype}\n")
+        if lora_config:
+            f.write(f"\nLoRA Configuration:\n")
+            f.write(f"  Rank (r): {lora_config.r}\n")
+            f.write(f"  Alpha: {lora_config.lora_alpha}\n")
+            f.write(f"  Dropout: {lora_config.lora_dropout}\n")
+            f.write(f"  Target modules: {lora_config.target_modules}\n")
+    
+    print(f"✓ 合并信息保存到: {info_path}")
+    
+    return merged_model
+
+
+def merge_and_save_full_pipeline(
+    base_model_path: str,
+    lora_weights_path: str,
+    output_path: str,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str = "cpu"
+):
+    """
+    合并 LoRA 权重并保存完整的 FluxFillPipeline
+    
+    Args:
+        base_model_path: 基础 Flux Fill 模型路径
+        lora_weights_path: PEFT LoRA 权重路径
+        output_path: 输出路径（保存完整 pipeline）
+        dtype: 数据类型
+        device: 加载设备
+    """
+    print("=" * 80)
+    print("Step 1: 加载基础 Flux Fill Pipeline...")
+    print("=" * 80)
+    
+    # 加载完整 pipeline
+    pipe = FluxFillPipeline.from_pretrained(
+        base_model_path,
+        torch_dtype=dtype,
+        low_cpu_mem_usage=True
+    )
+    
+    print(f"✓ Pipeline 加载完成: {base_model_path}")
+    
+    # 移动到指定设备
+    if device != "cpu":
+        print(f"  移动 transformer 到 {device}...")
+        pipe.transformer = pipe.transformer.to(device)
+    
+    print("\n" + "=" * 80)
+    print("Step 2: 加载并合并 PEFT LoRA 权重...")
+    print("=" * 80)
+    
+    # 加载 PEFT 模型
+    peft_model = PeftModel.from_pretrained(
+        pipe.transformer,
+        lora_weights_path,
+        is_trainable=False
+    )
+    peft_model.set_adapter("default")
+    
+    # 合并权重
+    merged_transformer = peft_model.merge_and_unload()
+    
+    # 替换 pipeline 中的 transformer
+    pipe.transformer = merged_transformer
+    
+    print(f"✓ 权重合并完成")
+    
+    print("\n" + "=" * 80)
+    print("Step 3: 保存完整 Pipeline...")
+    print("=" * 80)
+    
+    # 创建输出目录
+    os.makedirs(output_path, exist_ok=True)
+    
+    # 保存完整 pipeline
+    print(f"  保存到 {output_path}...")
+    pipe.save_pretrained(
+        output_path,
+        safe_serialization=True,
+        max_shard_size="5GB"
+    )
+    
+    print(f"✓ 完整 Pipeline 保存完成: {output_path}")
+    
+    # 保存合并信息
+    info_path = os.path.join(output_path, "merge_info.txt")
+    with open(info_path, "w") as f:
+        f.write(f"Base model: {base_model_path}\n")
+        f.write(f"LoRA weights: {lora_weights_path}\n")
+        f.write(f"Merged pipeline: {output_path}\n")
+        f.write(f"Data type: {dtype}\n")
+        f.write(f"Components saved:\n")
+        f.write(f"  - Transformer (merged with LoRA)\n")
+        f.write(f"  - VAE\n")
+        f.write(f"  - Text Encoder\n")
+        f.write(f"  - Scheduler\n")
+        f.write(f"  - Other components\n")
+    
+    print(f"✓ 合并信息保存到: {info_path}")
+    
+    return pipe
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="将 PEFT LoRA 权重合并到基础 Flux Transformer 模型中"
+    )
+    
+    parser.add_argument(
+        "--base_model_path",
+        type=str,
+        default="/home/tione/notebook/research/retofan/ckpt/FLUX.1-Fill-dev-UDG-1121_e4",
+        help="基础 Flux Fill 模型路径"
+    )
+    
+    parser.add_argument(
+        "--lora_weights_path",
+        type=str,
+        default="/home/tione/notebook2/research/retofan/code/RL/flow_grpo/logs/defectgen_det/flux_fill_redux/checkpoints/checkpoint-60/lora",
+        help="PEFT LoRA 权重路径"
+    )
+    
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="/home/tione/notebook/research/retofan/ckpt/FLUX.1-Fill-dev-UDG-1121_e4_defect_gen_det_e60",
+        help="输出路径（保存合并后的模型）"
+    )
+    
+    parser.add_argument(
+        "--save_full_pipeline",
+        action="store_true",
+        help="保存完整 FluxFillPipeline（包含 VAE、Text Encoder 等），而不只是 Transformer"
+    )
+    
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["float32", "float16", "bfloat16"],
+        help="数据类型"
+    )
+    
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="加载设备（cpu 或 cuda:0 等）。建议使用 cpu 以节省显存"
+    )
+    
+    args = parser.parse_args()
+    
+    # 转换 dtype
+    dtype_map = {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16
+    }
+    dtype = dtype_map[args.dtype]
+    
+    # 检查路径
+    if not os.path.exists(args.base_model_path):
+        print(f"错误: 基础模型路径不存在: {args.base_model_path}")
+        return
+    
+    if not os.path.exists(args.lora_weights_path):
+        print(f"错误: LoRA 权重路径不存在: {args.lora_weights_path}")
+        return
+    
+    print("\n" + "=" * 80)
+    print("PEFT LoRA 权重合并工具")
+    print("=" * 80)
+    print(f"基础模型: {args.base_model_path}")
+    print(f"LoRA 权重: {args.lora_weights_path}")
+    print(f"输出路径: {args.output_path}")
+    print(f"保存类型: {'完整 Pipeline' if args.save_full_pipeline else '仅 Transformer'}")
+    print(f"数据类型: {args.dtype}")
+    print(f"加载设备: {args.device}")
+    print("=" * 80 + "\n")
+    
+    try:
+        if args.save_full_pipeline:
+            # 保存完整 pipeline
+            merge_and_save_full_pipeline(
+                base_model_path=args.base_model_path,
+                lora_weights_path=args.lora_weights_path,
+                output_path=args.output_path,
+                dtype=dtype,
+                device=args.device
+            )
+        else:
+            # 只保存 transformer
+            merge_and_save_transformer(
+                base_model_path=args.base_model_path,
+                lora_weights_path=args.lora_weights_path,
+                output_path=args.output_path,
+                dtype=dtype,
+                device=args.device
+            )
+        
+        print("\n" + "=" * 80)
+        print("✅ 合并完成！")
+        print("=" * 80)
+        print(f"\n合并后的模型已保存到: {args.output_path}")
+        print("\n使用方法:")
+        
+        if args.save_full_pipeline:
+            print("  # 直接加载合并后的完整 pipeline")
+            print("  from diffusers import FluxFillPipeline")
+            print(f"  pipe = FluxFillPipeline.from_pretrained('{args.output_path}')")
+        else:
+            print("  # 加载基础 pipeline，然后替换 transformer")
+            print("  from diffusers import FluxFillPipeline")
+            print("  from diffusers.models import FluxTransformer2DModel")
+            print(f"  pipe = FluxFillPipeline.from_pretrained('{args.base_model_path}')")
+            print(f"  pipe.transformer = FluxTransformer2DModel.from_pretrained('{args.output_path}')")
+        
+        print("\n" + "=" * 80)
+        
+    except Exception as e:
+        print(f"\n❌ 错误: {e}")
+        import traceback
+        traceback.print_exc()
+        return 1
+    
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())