# PortraitCraft Challenge Track 2: Portrait Composition Generation - Top 1 Solution

This repository contains the Top 1 solution code for the **PortraitCraft Challenge Track 2: Portrait Composition Generation**.

Our solution leverages a powerful diffusion pipeline built upon Z-Image-Turbo, enhanced with ControlNet for precise pose conditioning and fine-tuned using DPO-LoRA.

***

## 🚀 Getting Started

### 1. Model Preparation

Before running the code, you need to download the required base models and ControlNet models into the `models` directory. You can choose to download them from either Hugging Face or ModelScope.

**From Hugging Face:**

- `Tongyi-MAI/Z-Image-Turbo`
- `alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1`

**From ModelScope:**

- `Tongyi-MAI/Z-Image-Turbo`
- `PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1`

Ensure the downloaded weights are placed correctly so the scripts can load them properly.

***

### 2. Inference

To generate portrait images based on the provided JSON tasks (which include text prompts and control pose prompt), run the evaluation script:

```bash
python inference/eval.py
```

The script will load the base model, ControlNet, and the trained LoRA weights to generate the  final submitted images.

***

### 3. Calculate Parameters

If you want to check the parameter count of the complete model pipeline (including DiT, Text Encoder, VAE, ControlNet, and LoRA), run:

```bash
python inference/calculate_params.py
```

**Expected Output:**

```text
=== Components Parameter Count ===
DiT (Transformer): 6.1549 B
Text Encoder: 4.0225 B
VAE (Encoder + Decoder): 0.0838 B
ControlNet: 3.3562 B
LoRA (Included in DiT): 0.1588 B

=== Parameter Summary ===
Total Model Parameters: 13.6174 B
Total Model Parameters: 13,617,422,307
```

*(Note: During inference, the LoRA weights are fused directly into the DiT transformer backbone to avoid extra computation overhead. Thus, the total parameter count remains highly efficient.)*

***

### 4. Training

We use a DPO-based LoRA training strategy combined with ControlNet to align the model with human preferences and improve composition quality.

To start the training process, execute the provided bash script:

```bash
bash train/train_ai4va_pose_controlnet_dpo_lora_8x1.sh
```