retofan23333/UniDG-SFT-LoRA-Release

UniDG-SFT-LoRA

LoRA weights for UniDG (Universal Defect Generation), trained via Diversity-SFT with complementary sampling on the UDG dataset (300K quadruplets).

[Paper] [Code] [UniDG-RFT-LoRA]

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 Diversity-SFT variant, optimized for diverse defect generation.

Variant	Training	Focus
UniDG-SFT (this)	Diversity-SFT with complementary sampling	Diverse defect patterns
UniDG-RFT	Consistency-RFT with Flow-GRPO + dual rewards	Consistent & faithful defects

Usage

Requirements

• FLUX.1-Fill-dev (inpainting backbone)
• FLUX.1-Redux-dev (reference conditioning)

Quick Start

from unidg import ImageUniDG
from PIL import Image
import torch

model = ImageUniDG(
    flux_model_path="path/to/FLUX.1-Fill-dev",
    redux_model_path="path/to/FLUX.1-Redux-dev",
    lora_weights_path="path/to/UniDG-SFT-LoRA/lora_weights.safetensors",
    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")

See the inference repo for Gradio demo and full documentation.

Citation

@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}
}