tsiendragon/qwen-image-finetune/qwen-image-edit-aipc-finetune-08-validation

validation:
  enabled: true
  steps: 100          # generate samples every 100 training steps
  max_samples: 2      # 2-5 is practical; more = slower validation pass
  seed: 42            # fixed seed for consistent comparison across checkpoints
  dataset:
    class_path: "qflux.data.dataset.ImageDataset"
    init_args:
      dataset_path:
        - split: test
          repo_id: <your-dataset>
      selected_control_indexes: [1]
      processor:
        class_path: "qflux.data.preprocess.ImageProcessor"
        init_args:
          process_type: center_crop
          target_size: [384, 672]
          controls_size: [[384, 672], [512, 512]]

tensorboard --logdir=outputs/<run_name>/logs

# inference_compare.py
# Run inference on a test split with and without LoRA for visual comparison.
# Supported test data sources (pass exactly one):
#   --test-parquet  parquet file or directory containing data/*.parquet files
#   --test-dir      local directory with training_images/ + control_images/ layout
#   --test-csv      CSV with columns: path_target, path_control[_N], prompt
# Usage:
#   python inference_compare.py --config config.yaml --test-dir <test-folder>
#       --output-dir outputs/compare/<run>/base --num-samples 5
#   python inference_compare.py --config config.yaml --test-dir <test-folder>
#       --output-dir outputs/compare/<run>/lora --num-samples 5
#       --lora-weight outputs/<run>/checkpoint-<step>/pytorch_lora_weights.safetensors
# Output per dir: sample_NN_<id>.png, sample_NN_<id>__target.png, manifest.json
# Prerequisite: run from the qwen-image-finetune project root with qflux importable (src/ on PYTHONPATH).
import argparse, csv, io, json, os, re, sys
from pathlib import Path
import PIL.Image
import torch
from qflux.data.config import Config, TrMode, load_config_from_yaml
try:
    from qflux.utils.memory_probe import MemoryProbe
except ImportError:
    MemoryProbe = None
_IMG_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
_IMG_DIRS = ["training_images", "images", "target_images", "target", "targets"]
_CTRL_DIRS = ["control_images", "control", "condition_images", "controls"]
_CTRL_EXTRA = re.compile(r"_control_(\d+)\.[a-zA-Z]+$")
_MASK_PAT = re.compile(r"_mask\.[a-zA-Z]+$")
def _open_image(path):
    return PIL.Image.open(path).convert("RGB")
def _read_parquet_samples(parquet_path, n):
    try:
        import pyarrow.parquet as pq
    except ImportError:
        raise SystemExit("pyarrow is required for --test-parquet. Install: pip install pyarrow")
    if parquet_path.is_dir():
        candidates = sorted(parquet_path.glob("data/test-*.parquet"))
        if not candidates:
            candidates = sorted(parquet_path.glob("data/*.parquet"))
        if not candidates:
            raise SystemExit(f"No parquet files found under {parquet_path}/data/.")
        table = pq.read_table(candidates[0].as_posix())
    else:
        table = pq.read_table(parquet_path.as_posix())
    rows = table.slice(0, n).to_pylist()
    samples = []
    for i, row in enumerate(rows):
        ctrl_imgs = []
        for entry in row.get("control_images") or []:
            if isinstance(entry, dict) and "bytes" in entry and entry["bytes"]:
                ctrl_imgs.append(PIL.Image.open(io.BytesIO(entry["bytes"])).convert("RGB"))
            elif isinstance(entry, (bytes, bytearray)):
                ctrl_imgs.append(PIL.Image.open(io.BytesIO(entry)).convert("RGB"))
        target_entry = row.get("target_image")
        tgt = None
        if isinstance(target_entry, dict) and target_entry.get("bytes"):
            tgt = PIL.Image.open(io.BytesIO(target_entry["bytes"])).convert("RGB")
        samples.append({"id": row.get("id", str(i)), "control_images": ctrl_imgs,
                        "target_image": tgt, "prompt": row.get("prompt", "")})
    return samples
def _read_local_dir_samples(dataset_dir, n):
    images_dir = next((dataset_dir / d for d in _IMG_DIRS if (dataset_dir / d).exists()), None)
    control_dir = next((dataset_dir / d for d in _CTRL_DIRS if (dataset_dir / d).exists()), None)
    if images_dir is None or control_dir is None:
        raise SystemExit(f"Could not find image/control subdirectories in {dataset_dir}.")
    stems = [f.stem for f in sorted(images_dir.iterdir())
             if f.suffix.lower() in _IMG_EXTS
             and not _MASK_PAT.search(f.name) and not _CTRL_EXTRA.search(f.name)]
    samples = []
    for stem in stems:
        if len(samples) >= n:
            break
        target_path = next((images_dir / f"{stem}{e}" for e in _IMG_EXTS
                             if (images_dir / f"{stem}{e}").exists()), None)
        main_ctrl = next((control_dir / f"{stem}{e}" for e in _IMG_EXTS
                          if (control_dir / f"{stem}{e}").exists()), None)
        if target_path is None or main_ctrl is None:
            continue
        extra_ctrls, idx = [], 1
        while True:
            ec = next((control_dir / f"{stem}_control_{idx}{e}" for e in _IMG_EXTS
                       if (control_dir / f"{stem}_control_{idx}{e}").exists()), None)
            if ec is None:
                break
            extra_ctrls.append(ec); idx += 1
        prompt_path = (images_dir / f"{stem}.txt") if (images_dir / f"{stem}.txt").exists() \
            else (control_dir / f"{stem}.txt") if (control_dir / f"{stem}.txt").exists() else None
        prompt = prompt_path.read_text(encoding="utf-8").strip() if prompt_path else ""
        ctrl_imgs = [_open_image(str(main_ctrl))] + [_open_image(str(p)) for p in extra_ctrls]
        samples.append({"id": stem, "control_images": ctrl_imgs,
                        "target_image": _open_image(str(target_path)), "prompt": prompt})
    if not samples:
        raise SystemExit(f"No valid samples found in {dataset_dir}.")
    return samples
def _read_csv_samples(csv_path, n):
    with open(csv_path, encoding="utf-8", newline="") as f:
        rows = list(csv.DictReader(f))
    if not rows:
        raise SystemExit(f"CSV {csv_path} is empty.")
    control_cols = sorted(c for c in rows[0] if c.startswith("path_control"))
    if not control_cols:
        raise SystemExit(f"CSV {csv_path} has no path_control* columns.")
    samples = []
    for i, row in enumerate(rows[:n]):
        ctrl_imgs = [_open_image(row[c]) for c in control_cols
                     if row.get(c) and os.path.exists(row[c])]
        if not ctrl_imgs:
            continue
        tgt_path = row.get("path_target", "")
        tgt = _open_image(tgt_path) if tgt_path and os.path.exists(tgt_path) else None
        samples.append({"id": row.get("id", str(i)), "control_images": ctrl_imgs,
                        "target_image": tgt, "prompt": row.get("prompt", "")})
    if not samples:
        raise SystemExit(f"No valid samples loaded from CSV {csv_path}.")
    return samples
def _load_test_samples(args, n):
    if args.test_parquet:
        return _read_parquet_samples(Path(args.test_parquet), n)
    if args.test_dir:
        return _read_local_dir_samples(Path(args.test_dir), n)
    return _read_csv_samples(Path(args.test_csv), n)
def _instantiate_trainer(config):
    if config.trainer_type == "QwenImageEdit":
        from qflux.trainer.qwen_image_edit_trainer import QwenImageEditTrainer
        return QwenImageEditTrainer(config)
    if config.trainer_type == "QwenImageEditPlus":
        from qflux.trainer.qwen_image_edit_plus_trainer import QwenImageEditPlusTrainer
        return QwenImageEditPlusTrainer(config)
    raise NotImplementedError(f"supports QwenImageEdit and QwenImageEditPlus, got {config.trainer_type}")
def main():
    p = argparse.ArgumentParser(description="Run inference on a test split for visual LoRA comparison.")
    p.add_argument("--config", required=True)
    src = p.add_mutually_exclusive_group(required=True)
    src.add_argument("--test-parquet",
                     help="Parquet file or directory containing data/*.parquet files")
    src.add_argument("--test-dir",
                     help="Local directory with training_images/ + control_images/ layout")
    src.add_argument("--test-csv",
                     help="CSV with columns: path_target, path_control[_N], prompt")
    p.add_argument("--output-dir", required=True)
    p.add_argument("--num-samples", type=int, default=5)
    p.add_argument("--num-inference-steps", type=int, default=20)
    p.add_argument("--lora-weight", default=None)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--text-encoder-device", choices=["cpu", "xpu:0"], default="cpu",
                   help="text_encoder placement; default 'cpu' (Config B, bf16 on CPU). Use 'xpu:0' "
                        "for Config A — requires --text-encoder-path and a raised Step 1 §1.3 XPU cap.")
    p.add_argument("--text-encoder-path", default=None,
                   help="Dir of a pre-quantized NF4 text_encoder (Step 6 §6.2); overrides "
                        "model.text_encoder_path. Required for Config A (--text-encoder-device xpu:0).")
    args = p.parse_args()