HW6 Brainrot Image Generation

基本資訊

本作業的目標是訓練一個 conditional image generation 模型，根據指定的 animal-object pair 生成 64×64 的 Brainrot 圖像。模型需要依照給定條件產生對應圖片，並透過 FID 評估生成影像的品質與分布相似度，以及透過 CLIP-T 評估生成圖片與文字條件之間的語意一致性。

實作概要

• Pixel-space diffusion
• DDPM training objective
• Epsilon prediction
• DDIM sampler
• Classifier-Free Guidance
• CFG dropout
• EMA
• Min-SNR loss weighting
• CLIP text embedding、Animal embedding、Object embedding
• CLIP-T data filtering
• Pretrained model: ViT-B-32-quickgelu, pretrained="openai"

實作結果

FID	CLIP-T
41.91	0.2706

產生此結果的 checkpoint 可以透過以下指令下載:

hf download jimmy0214/GAI-HW6-checkpoint \
  --repo-type model \
  --local-dir ./checkpoint

使用前請先確認已安裝 huggingface_hub

# 安裝 huggingface_hub
pip install -U huggingface_hub
# 登入 huggingface
hf auth login

環境安裝

請先進入專案根目錄：

cd GAI-HW6-Brainrot-Image-Generation/

Step 1: 建立 conda 環境

conda create -n GAI_HW6 python=3.11 -y
conda activate GAI_HW6

Step 2: 安裝相依套件

pip install -r requirements.txt

資料夾結構

GAI-HW6-Brainrot-Image-Generation/
├── scripts/
│   ├── brainrot_ddpm.py       # Dataset、模型、diffusion、sampling
│   ├── clip_t_audit.py        # CLIP-T 資料清理
│   ├── clip_t_class_stats.py  # 用於分析生成圖片的 CLIP-T 分數
│   ├── score.py               # 評分程式碼
│   ├── train.py               # 訓練程式
│   └── inference.py           # 推論程式
├── requirements.txt
├── CLI.md
└── README.md

clip_t_class_stats.py 與 score.py 的使用方式請參閱 CLI.md 中的 score.py 使用方式。

使用須知

請先準備好以下檔案或資料夾:

├── trainset/         # 訓練圖片
│   ├── 000001.png
│   ├── 000002.png
│   └── 000003.png
├── train.csv         # 記錄每張訓練圖片的 id、對應的 animal 與 object 種類
└── generate.csv      # 紀錄每張待生成圖片的 id、對應的 condition

或是透過以下指令下載，其中也包含使用 score.py 所需的資料夾:

hf download jimmy0214/GAI-HW6-dataset \
  --repo-type dataset \
  --local-dir .

流程概述

先使用 clip_t_audit.py 計算訓練圖片與 prompt 的 CLIP-T 分數，紀錄部分 CLIP-T 分數低的資料在 JSON 檔案中，用於在訓練時忽略這些資料；接著使用 train.py 讀取該 JSON 進行訓練；最後使用 inference.py 載入 checkpoint，並根據 generate.csv 生成圖片。

資料清理：scripts/clip_t_audit.py

此指令會計算訓練圖片與 prompt 的 CLIP-T 分數，並針對每個 (animal, object) 類別產生排除清單，以下設定將每個類別最低的十筆資料排除同時每個類別最少有十筆資料。

python3 scripts/clip_t_audit.py \
  --train_csv path/to/train.csv \
  --image_dir path/to/trainset \
  --output_json train_clip_t_scores.json \
  --lowest_json train_clip_t_lowest_by_class.json \
  --exclude_json train_clip_t_exclude_10.json \
  --lowest_k 10 \
  --drop_lowest_k 10 \
  --min_remaining_per_class 10 \
  --class_mode pair \
  --batch_size 64 \
  --device cuda

輸出說明：

├── train_clip_t_scores.json            # 記錄每張圖片的 CLIP-T 分數
├── train_clip_t_lowest_by_class.json   # 紀錄每個類別前幾低分的資料
└── train_clip_t_exclude_10.json        # 排除清單，後續使用的即為此檔案

訓練：scripts/train.py

此指令會使用 CLIP-T 排除清單訓練 pixel-space diffusion model，啟用 cosine beta schedule、Min-SNR loss weighting 與 EMA。

python3 scripts/train.py \
  --image_dir path/to/trainset \
  --train_csv path/to/train.csv \
  --exclude_json train_clip_t_exclude_10.json \
  --output_dir result \
  --batch_size 16 \
  --num_epochs 400 \
  --beta_schedule cosine \
  --min_snr_gamma 5.0 \
  --guidance_scale 1.5 \
  --save_every_epochs 50 \
  --sample_every_epochs 50

訓練輸出會產生在 result/，輸出如下：

result/
├── best_model.pth              # loss 最低的 checkpoint
├── checkpoints/
│   ├── model_epoch_0050.pth
│   ├── model_epoch_0100.pth
│   └── model_epoch_0150.pth
├── latest_checkpoint.txt       # 紀錄最後儲存的 checkpoint，此檔案用於從中途開始訓練
├── samples/                    
│   ├── samples_epoch_0050.png
│   ├── samples_epoch_0100.png
│   └── samples_epoch_0150.png
├── training_history.json
└── training_curves.png

推論：scripts/inference.py

此指令會載入訓練好的 checkpoint，使用 DDIM 產生圖片；若 checkpoint 內含 EMA 權重，預設會使用 EMA。

python3 scripts/inference.py \
  --checkpoint result/best_model.pth \
  --generate_csv path/to/generate.csv \
  --output_dir res \
  --batch_size 128 \
  --num_sample_steps 100 \
  --sampler ddim \
  --ddim_eta 0.0 \
  --guidance_scale 3.5

產生的 2000 張圖片會產生在 res/ 資料夾下。