UniFit: Towards Universal Virtual Try-on with MLLM-Guided Semantic Alignment

This repository provides Flux.1 Fill inference for UniFit.

📢 News

• [2025-11-20] 🎉 UniFit has been accepted to AAAI 2026!
• [2025-11-20] 🚀 Flux.1 Fill inference code is released.

📝 To-Do List

• Release Paper (arXiv)
• Release Flux.1 Fill Inference Code
• Release UniFit-Flux Pretrained Model (Checkpoint)
• Release Training Codes
  • Data processing scripts
  • Training scripts for Stage I & II

💡 Abstract

Image-based virtual try-on (VTON) aims to synthesize photorealistic images of a person wearing specified garments. Despite significant progress, building a universal VTON framework that can flexibly handle diverse and complex tasks remains a major challenge. Recent methods explore multi-task VTON frameworks guided by textual instructions, yet they still face two key limitations: (1) semantic gap between text instructions and reference images, and (2) data scarcity in complex scenarios. To address these challenges, we propose UniFit, a universal VTON framework driven by a Multimodal Large Language Model (MLLM). Specifically, we introduce an MLLM-Guided Semantic Alignment Module (MGSA), which integrates multimodal inputs using an MLLM and a set of learnable queries. By imposing a semantic alignment loss, MGSA captures cross-modal semantic relationships and provides coherent and explicit semantic guidance for the generative process, thereby reducing the semantic gap. Moreover, by devising a two-stage progressive training strategy with a self-synthesis pipeline, UniFit is able to learn complex tasks from limited data. Extensive experiments show that UniFit not only supports a wide range of VTON tasks, including multi-garment and model-to-model try-on, but also achieves state-of-the-art performance.

🔧 Installation

This project uses a uv-managed Python environment. In the current checkout, the environment is located at /root/data/Method/UniFit/paper/flux/FluxFill_v2_1024/code/UniFit/.venv; the commands below use the repository-relative path .venv so the instructions remain portable.

git clone https://github.com/zwplus/UniFit.git
cd UniFit

# Create the environment only if .venv does not already exist
uv venv --python 3.10 .venv

# Install all dependencies into this exact environment
uv pip install --python .venv/bin/python -r requirements.txt

# Optional: activate it for interactive use
source .venv/bin/activate

When the environment already exists, skip uv venv and run the uv pip install command to synchronize dependencies. .venv/ is a local environment and is not part of the repository.

⬇️ Download Model

Download the three model directories:

# Log in first if the FLUX model asks for authentication or license acceptance
hf auth login

# FLUX.1-Fill-dev base model
hf download black-forest-labs/FLUX.1-Fill-dev --local-dir ./ckpt/FLUX

# UniFit trained checkpoint
hf download zwpro/UniFit --local-dir ./ckpt/UniFit

# Official Qwen2-VL processor/model files
hf download Qwen/Qwen2-VL-2B-Instruct --local-dir ./ckpt/Qwen2-VL-2B-Instruct

The FLUX model may require Hugging Face authentication and acceptance of its model terms. The UniFit checkpoint directory must directly contain pose_net/, qwenvl/, qwenvl_connector/, qwenvl_learn_query/, qwenvl_pool_connector/, and transformer/. The runtime loads the trained Qwen2-VL weights from ./ckpt/UniFit/qwenvl; the official Qwen2-VL directory supplies the processor files used by --qwenvl_base_path.

🚀 Inference

A single script unified_flux_inference.py covers all 6 try-on tasks.

.venv/bin/python unified_flux_inference.py \
  --task_type <task_type> \
  --prompt <prompt> \
  --output_dir <output_dir> \
  --flux_path <flux_fill_path> \
  --ckpt_path <unifit_ckpt_path> \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct \
  --height 1024 \
  --width 768 \
  [task-specific image paths]

The UniFit Qwen2-VL model used for inference is loaded from <unifit_ckpt_path>/qwenvl. The standard Qwen2-VL path above supplies only the tokenizer, image processor, and chat template. It can also be a local directory containing the official Qwen2-VL-2B-Instruct processor files.

Task types

--task_type	Required inputs	Default output resolution
tryon	--model_path, --cloth_paths, --pose_path	1024×768 (configurable)
cloth restruct	--model_path	1024×768 (configurable)
model-free tryon	--cloth_paths	1024×768 (configurable)
model-to-model tryon	--model_path, --cloth_paths, --pose_path	1024×768 (configurable)
multi-view tryon	--model_path, --cloth_paths (front, back; extra entries ignored), --pose_path	1024×768 (configurable)
multi-cloth tryon	--model_path, --cloth_paths (upper, lower; extra entries ignored), --pose_path	1024×768 (configurable)

Examples

Try-on

.venv/bin/python unified_flux_inference.py \
  --task_type tryon \
  --prompt "Replace the black t-shirt with leopard print with the dark red t-shirt, which features a white graphic." \
  --model_path ./assets/input/tryon/model.jpg \
  --cloth_paths ./assets/input/tryon/cloth.jpg \
  --pose_path ./assets/input/tryon/dwpose.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Cloth reconstruction

.venv/bin/python unified_flux_inference.py \
  --task_type "cloth restruct" \
  --prompt "Reconstruct the grey short-sleeved t-shirt with the rectangular red, white, and blue logo that the model is wearing." \
  --model_path ./assets/input/cloth_restruct/model.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Model-free try-on

.venv/bin/python unified_flux_inference.py \
  --task_type "model-free tryon" \
  --prompt "a young adult female wearing a navy blue cotton crew-neck t-shirt with white ribbed trim around the neckline and sleeve openings, featuring three distinct blue stripes running down each shoulder to the sleeve opening and a small blue Adidas trefoil logo printed on the left chest, along with bright blue relaxed-fit track pants with white stripes down the sides and a white Adidas logo near the cuff. The young adult female with dark hair pulled back and large silver hoop earrings, should be standing with her weight slightly shifted to one side, looking directly at the camera with a serious expression, and her hands hanging naturally at her sides. This t-shirt should be worn untucked, appearing to have a slightly relaxed yet fitted cut, ending at the upper hip, with the crew neckline and sleeve openings showing the white ribbed trim clearly and the short sleeves ending mid-bicep. The background should be a plain white studio backdrop.a young adult female wearing a navy blue cotton crew-neck t-shirt with white ribbed trim around the neckline and sleeve openings, featuring three distinct blue stripes running down each shoulder to the sleeve opening and a small blue Adidas trefoil logo printed on the left chest, along with bright blue relaxed-fit track pants with white stripes down the sides and a white Adidas logo near the cuff. The young adult female with large silver hoop earrings, should be standing with her weight slightly shifted to one side, looking directly at the camera with a serious expression, and her hands hanging naturally at her sides. This t-shirt should be worn untucked, appearing to have a slightly relaxed yet fitted cut, ending at the upper hip, with the crew neckline and sleeve openings showing the white ribbed trim clearly and the short sleeves ending mid-bicep. The background should be a plain white studio backdrop. " \
  --cloth_paths ./assets/input/model-free_tryon/cloth.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Model-to-model try-on

.venv/bin/python unified_flux_inference.py \
  --task_type "model-to-model tryon" \
  --prompt "Replace the white long-sleeved collared shirt on the original model with the light blue sleeveless ruffled knit top from the clothing mannequin." \
  --model_path ./assets/input/model-to-model_tryon/model.jpg \
  --cloth_paths ./assets/input/model-to-model_tryon/cloth.jpg \
  --pose_path ./assets/input/model-to-model_tryon/dwpose.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Multi-view try-on

Front View

.venv/bin/python unified_flux_inference.py \
  --task_type "multi-view tryon" \
  --prompt "Replace the white long-sleeved collared shirt on the original model with the light blue sleeveless ruffled knit top from the clothing mannequin." \
  --model_path ./assets/input/multi-view_tryon/model_front.jpg \
  --cloth_paths ./assets/input/multi-view_tryon/cloth_front.jpg ./assets/input/multi-view_tryon/cloth_back.jpg \
  --pose_path ./assets/input/multi-view_tryon/dwpose_front.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Back View

.venv/bin/python unified_flux_inference.py \
  --task_type "multi-view tryon" \
  --prompt "Replace the white long-sleeved collared shirt on the original model with the light blue sleeveless ruffled knit top from the clothing mannequin." \
  --model_path ./assets/input/multi-view_tryon/model_back.jpg \
  --cloth_paths ./assets/input/multi-view_tryon/cloth_front.jpg ./assets/input/multi-view_tryon/cloth_back.jpg \
  --pose_path ./assets/input/multi-view_tryon/dwpose_back.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Multi-cloth try-on

.venv/bin/python unified_flux_inference.py \
  --task_type "multi-cloth tryon" \
  --prompt "Replace the model's white strappy top with the black and white color-block shirt with a large lapel collar from image one, and replace the black trousers with the blue jeans from image two." \
  --model_path ./assets/input/multi-cloth_tryon/model.jpg \
  --cloth_paths ./assets/input/multi-cloth_tryon/upper_cloth.jpg ./assets/input/multi-cloth_tryon/lower_cloth.jpg \
  --pose_path ./assets/input/multi-cloth_tryon/dwpose.jpg \
  --output_dir ./output \
  --flux_path ./ckpt/FLUX \
  --ckpt_path ./ckpt/UniFit \
  --qwenvl_base_path ./ckpt/Qwen2-VL-2B-Instruct

Results are saved as:

output_dir/
└── <task_type>/
    └── <image_name>/
        └── output.png

📂 Code Structure

UniFit/
├── unified_flux_inference.py   # unified inference entry point for all 6 tasks
├── utils.py                    # Qwen2Connector and shared helpers
├── flux/                       # Flux pipeline, transformer, attention, and PoseNet
│   ├── flux_qwen_pipelines.py
│   ├── transformer_flux.py
│   ├── attention_processor.py
│   └── pose_net.py
├── qwen2_vl/                   # UniFit's Qwen2-VL and MGSA implementation
│   ├── modeling_qwen2_vl.py
│   ├── processing_qwen2_vl.py
│   ├── image_processing_qwen2_vl.py
│   └── configuration_qwen2_vl.py
├── assets/
│   ├── input/                  # small local input examples for the README commands
│   └── network.png             # method overview figure
├── ckpt/                      # local downloaded FLUX, UniFit, and Qwen2-VL weights
├── .venv/                     # local uv-managed Python environment (ignored by Git)
└── requirements.txt           # Python dependencies installed with uv

🙏 Acknowledgement

Our code is modified based on Diffusers. We use FLUX.1-Fill-dev as the base model, and Qwen2-VL-2B-Instruct as the MGSA Module. Thanks to all the contributors!