tsiendragon/qwen-image-finetune/qwen-image-edit-aipc-finetune-04-config

# detect_trainer.py
# Detect which qflux trainer a model needs, from its model_index.json _class_name:
#   QwenImageEditPipeline      -> "QwenImageEdit"      (Qwen-Image-Edit)
#   QwenImageEditPlusPipeline  -> "QwenImageEditPlus"  (Qwen-Image-Edit-2509 / 2511)
# Usage:
#   python detect_trainer.py <model_dir>
#   python detect_trainer.py --json <model_dir>
# Exit codes: 0 = detected; 2 = model_index.json missing; 3 = unknown pipeline.
# Network policy: reads only a local model_index.json. A HuggingFace repo id (not a local dir)
# has no local file to read and exits 2 — download the model first (Step 1) or set trainer manually.
import argparse, json, os, sys
# Map a diffusers pipeline `_class_name` to the qflux trainer kind. Extend this table if
# upstream adds new Qwen-Image-Edit pipeline variants.
PIPELINE_TO_TRAINER = {
    "QwenImageEditPipeline": {"trainer": "QwenImageEdit", "multi_image": False,
                               "note": "Qwen-Image-Edit, single-image editing"},
    "QwenImageEditPlusPipeline": {"trainer": "QwenImageEditPlus", "multi_image": True,
                                   "note": "Qwen-Image-Edit-2509 / 2511, multi-image editing"},
}
def detect(model_dir):
    """Return {pipeline_class, trainer, multi_image, note} for a local model dir.
    Raises FileNotFoundError if model_index.json is absent, KeyError if the pipeline is unknown."""
    index_path = os.path.join(model_dir, "model_index.json")
    if not os.path.isfile(index_path):
        raise FileNotFoundError(index_path)
    with open(index_path, encoding="utf-8") as f:
        index = json.load(f)
    pipeline_class = index.get("_class_name")
    if pipeline_class not in PIPELINE_TO_TRAINER:
        raise KeyError(pipeline_class)
    entry = PIPELINE_TO_TRAINER[pipeline_class]
    return {"pipeline_class": pipeline_class, "trainer": entry["trainer"],
            "multi_image": entry["multi_image"], "note": entry["note"]}
def main():
    p = argparse.ArgumentParser(description="Detect the qflux trainer for a model from its model_index.json.")
    p.add_argument("model_dir", help="Local directory of the downloaded model")
    p.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
    args = p.parse_args()
    try:
        result = detect(args.model_dir)
    except FileNotFoundError as e:
        print(f"ERROR: {e} not found — pass a local model directory that contains model_index.json (Step 1).", file=sys.stderr)
        return 2
    except KeyError as e:
        known = ", ".join(sorted(PIPELINE_TO_TRAINER))
        print(f"ERROR: unrecognized pipeline _class_name={e!s} in {args.model_dir}/model_index.json. Known: {known}.", file=sys.stderr)
        return 3
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"{result['trainer']}   ({result['note']})")
    return 0
if __name__ == "__main__":
    sys.exit(main())

python detect_trainer.py <path\to\model>

python probe_hw.py > probe.json

# recommend_config.py
# Read probe JSON (from probe_hw.py) and emit a complete NF4-QLoRA training YAML that fits the
# machine. Trainer selection is MODEL-driven: pass --model-path so model_index.json picks the
# trainer. Without --model-path, trainer is a TODO placeholder.
# Resolution: when --dataset-path is given, target_size and controls_size are scaled to
# preserve the dataset's real aspect ratio at a conservative pixel budget (≈ 258 k–262 k px²).
# Without a dataset path, shape-bucket defaults are used. These are starting points; adjust
# based on your machine and confirm with monitor_xpu_memory.py (Step 7).
# Usage:
#   python probe_hw.py | python recommend_config.py
#   python recommend_config.py --probe probe.json --out config.yaml
#   python recommend_config.py --probe probe.json --model-path <model-dir> --ram-tier 32 --dataset-path <path> --out config.yaml
# Selected patterns are printed to stderr. Network policy: no network access.
import argparse
import glob
import io
import json
import os
import re
import sys
# ── Dataset introspection ──────────────────────────────────────────────────────
# Supports image directory, parquet, and CSV datasets.
# Returns the real target shape, control count, and per-control shapes from the
# first sample — so target_size and controls_size are seeded from actual data.
_IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
_IMG_DIRS = ["training_images", "images", "target_images", "target", "targets"]
_CTRL_DIRS = ["control_images", "control", "condition_images", "controls"]
def _detect_dataset_format(path):
    """Return 'image_directory' | 'parquet' | 'csv' | 'unknown'."""
    if not path:
        return "unknown"
    if isinstance(path, str) and path.lower().endswith(".csv"):
        return "csv"
    if os.path.isdir(path):
        if glob.glob(os.path.join(path, "data", "*.parquet")) or \
           glob.glob(os.path.join(path, "*.parquet")):
            return "parquet"
        return "image_directory"
    return "unknown"
def _img_hw_from_file(path):
    """(H, W) from an image file path using only stdlib struct (no PIL needed)."""
    try:
        import struct
        with open(path, "rb") as f:
            head = f.read(26)
        if head[:8] == b"\x89PNG\r\n\x1a\n":
            w, h = struct.unpack(">II", head[16:24])
            return (h, w)
        if head[:2] == b"\xff\xd8":
            with open(path, "rb") as f:
                f.read(2)
                while True:
                    marker = f.read(2)
                    if len(marker) < 2 or marker[0] != 0xFF:
                        break
                    seg_len = struct.unpack(">H", f.read(2))[0]
                    if marker[1] in (0xC0, 0xC2):
                        f.read(1)
                        h, w = struct.unpack(">HH", f.read(4))
                        return (h, w)
                    f.read(seg_len - 2)
    except Exception:
        pass
    return None
def _img_hw_from_cell(cell):
    """(H, W) from a parquet image cell (dict{bytes} / raw bytes). Requires PIL."""
    try:
        from PIL import Image
        if isinstance(cell, dict) and "bytes" in cell:
            cell = cell["bytes"]
        if isinstance(cell, (bytes, bytearray)):
            with Image.open(io.BytesIO(cell)) as im:
                w, h = im.size
            return (h, w)
        if hasattr(cell, "size"):
            w, h = cell.size
            return (h, w)
    except Exception:
        pass
    return None
def _introspect_local(path):
    img_dir = next((os.path.join(path, d) for d in _IMG_DIRS
                    if os.path.isdir(os.path.join(path, d))), None)
    ctrl_dir = next((os.path.join(path, d) for d in _CTRL_DIRS
                     if os.path.isdir(os.path.join(path, d))), None)
    if img_dir is None or ctrl_dir is None:
        return None
    ctrl_re = re.compile(r"_control_\d+\.(?:jpg|jpeg|png|bmp|webp)$", re.IGNORECASE)
    targets = [p for p in sorted(glob.glob(os.path.join(img_dir, "*.*")))
               if p.lower().endswith(_IMG_EXTS)
               and not os.path.basename(p).lower().endswith("_mask.png")
               and not ctrl_re.search(os.path.basename(p))]
    if not targets:
        return None
    stem = os.path.splitext(os.path.basename(targets[0]))[0]
    target_hw = _img_hw_from_file(targets[0])
    main_ctrl = next((os.path.join(ctrl_dir, stem + e) for e in _IMG_EXTS
                      if os.path.isfile(os.path.join(ctrl_dir, stem + e))), None)
    controls_hw = [_img_hw_from_file(main_ctrl)] if main_ctrl else []
    i = 1
    while True:
        ec = next((os.path.join(ctrl_dir, f"{stem}_control_{i}{e}") for e in _IMG_EXTS
                   if os.path.isfile(os.path.join(ctrl_dir, f"{stem}_control_{i}{e}"))), None)
        if ec is None:
            break
        controls_hw.append(_img_hw_from_file(ec))
        i += 1
    n = (1 if main_ctrl else 0) + (i - 1)
    return {"n_controls_total": n or None, "target_hw": target_hw, "controls_hw": controls_hw,
            "format": "image_directory"}
def _introspect_parquet(path):
    files = sorted(glob.glob(os.path.join(path, "data", "*.parquet")) or
                   glob.glob(os.path.join(path, "*.parquet")))
    train = [f for f in files if os.path.basename(f).startswith("train")] or files
    if not train:
        return None
    try:
        import pandas as pd
        df = pd.read_parquet(train[0])
    except Exception:
        return None
    if len(df) == 0:
        return None
    row = df.iloc[0]
    controls = row.get("control_images") if "control_images" in df.columns else None
    controls_hw = []
    n_total = 0
    if controls is not None:
        try:
            seq = list(controls)
        except TypeError:
            seq = [controls]
        n_total = len(seq)
        for c in seq:
            controls_hw.append(_img_hw_from_cell(c))
    target_hw = (_img_hw_from_cell(row.get("target_image"))
                 if "target_image" in df.columns else None)
    return {"n_controls_total": n_total or None, "target_hw": target_hw,
            "controls_hw": controls_hw, "format": "parquet"}
def _introspect_csv(path):
    try:
        import pandas as pd
        df = pd.read_csv(path)
    except Exception:
        return None
    if len(df) == 0:
        return None
    ctrl_cols = sorted(c for c in df.columns if "path_control" in c)
    row = df.iloc[0]
    base = os.path.dirname(path)
    def _resolve(rel):
        if not isinstance(rel, str):
            return None
        return rel if os.path.isabs(rel) else os.path.join(base, rel)
    controls_hw = []
    for c in ctrl_cols:
        p = _resolve(row.get(c))
        controls_hw.append(_img_hw_from_file(p) if p and os.path.isfile(p) else None)
    target_hw = None
    if "path_target" in df.columns:
        tp = _resolve(row.get("path_target"))
        if tp and os.path.isfile(tp):
            target_hw = _img_hw_from_file(tp)
    return {"n_controls_total": len(ctrl_cols) or None, "target_hw": target_hw,
            "controls_hw": controls_hw, "format": "csv"}
def introspect_dataset(dataset_path):
    """Introspect a local dataset and return shape/control info.