tonera

/

FLUX.2-klein-9B-Nunchaku

Like 18

I used this dataset: https://huggingface.co/datasets/UCSC-VLAA/HQ-Edit, and then extracted 200 images from it.

script code:

#!/usr/bin/env python3
"""
Export a bounded local calibration subset from supported image-edit datasets.

This script intentionally avoids `datasets.load_dataset(...)` because some
datasets are split across many very large parquet shards with embedded images.
Instead it:

1. Fetches the parquet shard manifest from the Hugging Face dataset server.
2. Downloads only the shard(s) needed for the requested sample count.
3. Extracts the source image + edit instruction rows locally from those shards.
4. Saves images and writes a DeepCompressor-compatible YAML prompt file.

Supported presets:
- `hq-edit`
- `magicbrush`
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import random
import shutil
import tempfile
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping


PARQUET_ENDPOINT = "https://datasets-server.huggingface.co/parquet"
DOWNLOAD_CHUNK_SIZE = 8 * 1024 * 1024
STATE_VERSION = 1
TMP_ROOT = Path(tempfile.gettempdir()) / "deepcompressor_hq_edit_export"
JOB_ROOT = TMP_ROOT / "jobs"
SHARD_CACHE_ROOT = TMP_ROOT / "shards"


@dataclass(frozen=True)
class DatasetPreset:
    repo_id: str
    image_column: str
    prompt_column: str
    extra_columns: tuple[str, ...] = ()
    default_split: str = "train"


DATASET_PRESETS: dict[str, DatasetPreset] = {
    "hq-edit": DatasetPreset(
        repo_id="UCSC-VLAA/HQ-Edit",
        image_column="input_image",
        prompt_column="edit",
    ),
    "magicbrush": DatasetPreset(
        repo_id="osunlp/MagicBrush",
        image_column="source_img",
        prompt_column="instruction",
        extra_columns=("turn_index",),
    ),
}


def log(message: str) -> None:
    print(f"[HQ-Edit Export] {message}", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download a bounded subset of an image-edit dataset and export a local DeepCompressor YAML prompt file."
    )
    parser.add_argument(
        "--dataset-preset",
        choices=sorted(DATASET_PRESETS.keys()),
        default="hq-edit",
        help="Built-in dataset preset. Default: %(default)s",
    )
    parser.add_argument(
        "--repo-id",
        default=None,
        help="Override the Hugging Face dataset repo id. By default this comes from --dataset-preset.",
    )
    parser.add_argument(
        "--split",
        default=None,
        help="Dataset split to use. By default this comes from --dataset-preset.",
    )
    parser.add_argument(
        "--sample-size",
        type=int,
        default=200,
        help="Number of samples to export. Must be > 0. Default: %(default)s",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Random seed used when shuffling shard order and row order. Default: %(default)s",
    )
    parser.add_argument(
        "--shuffle",
        action="store_true",
        help="Shuffle shard order and row order inside each downloaded shard before selecting samples.",
    )
    parser.add_argument(
        "--magicbrush-first-turn-only",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="For the magicbrush preset, keep only rows with turn_index == 1 so source images stay closer to real original photos. Default: %(default)s",
    )
    parser.add_argument(
        "--output-yaml",
        type=Path,
        required=True,
        help="Output YAML path, e.g. examples/diffusion/prompts/hq-edit-200.yaml",
    )
    parser.add_argument(
        "--image-dir",
        type=Path,
        required=True,
        help="Directory to save downloaded source images.",
    )
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=None,
        help="Directory to cache downloaded parquet shards. Default: a shared dataset cache under /tmp reused across runs.",
    )
    parser.add_argument(
        "--keep-shards",
        action="store_true",
        help="Keep downloaded parquet shards in the cache directory after export.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing YAML and image files if they already exist.",
    )
    parser.add_argument(
        "--reuse-existing-images",
        action="store_true",
        help="When an output image already exists, keep it and only rebuild the YAML mapping.",
    )
    return parser.parse_args()


def resolve_preset(args: argparse.Namespace) -> DatasetPreset:
    preset = DATASET_PRESETS[args.dataset_preset]
    repo_id = args.repo_id or preset.repo_id
    split = args.split or preset.default_split
    return DatasetPreset(
        repo_id=repo_id,
        image_column=preset.image_column,
        prompt_column=preset.prompt_column,
        extra_columns=preset.extra_columns,
        default_split=split,
    )


def build_job_config(
    args: argparse.Namespace,
    preset: DatasetPreset,
    output_yaml: Path,
    image_dir: Path,
) -> dict[str, Any]:
    return {
        "dataset_preset": args.dataset_preset,
        "repo_id": preset.repo_id,
        "split": preset.default_split,
        "sample_size": args.sample_size,
        "seed": args.seed,
        "shuffle": args.shuffle,
        "magicbrush_first_turn_only": args.magicbrush_first_turn_only,
        "output_yaml": output_yaml.as_posix(),
        "image_dir": image_dir.as_posix(),
    }


def resolve_resume_dir(job_config: Mapping[str, Any]) -> Path:
    digest = hashlib.sha256(json.dumps(job_config, sort_keys=True).encode("utf-8")).hexdigest()[:16]
    return (JOB_ROOT / digest).resolve()


def resolve_shard_cache_dir(preset: DatasetPreset, cache_dir: Path | None) -> Path:
    if cache_dir is not None:
        return cache_dir.resolve()
    shard_key = {
        "repo_id": preset.repo_id,
        "split": preset.default_split,
    }
    digest = hashlib.sha256(json.dumps(shard_key, sort_keys=True).encode("utf-8")).hexdigest()[:16]
    return (SHARD_CACHE_ROOT / digest).resolve()


def write_json_atomic(path: Path, payload: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(f"{path.name}.tmp")
    with tmp_path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2, sort_keys=True)
    os.replace(tmp_path, path)


def load_state(state_path: Path, job_config: Mapping[str, Any]) -> dict[str, Any]:
    if not state_path.exists():
        return {
            "version": STATE_VERSION,
            "job": dict(job_config),
            "exported": {},
            "shards": {},
        }
    with state_path.open("r", encoding="utf-8") as f:
        state = json.load(f)
    if state.get("version") != STATE_VERSION:
        raise SystemExit(f"Unsupported resume state version in {state_path}. Please remove the temp directory and retry.")
    if state.get("job") != dict(job_config):
        raise SystemExit(f"Resume state in {state_path} does not match the current command. Please remove the temp directory and retry.")
    state.setdefault("exported", {})
    state.setdefault("shards", {})
    return state


def save_state(state_path: Path, state: Mapping[str, Any]) -> None:
    write_json_atomic(state_path, state)


def advance_rng_for_shuffle(rng: random.Random, row_count: int) -> None:
    if row_count > 1:
        placeholder = list(range(row_count))
        rng.shuffle(placeholder)


def cleanup_temp_artifacts(resume_dir: Path) -> None:
    if resume_dir.exists():
        shutil.rmtree(resume_dir)


def import_dependencies():
    try:
        import yaml
    except ImportError as exc:  # pragma: no cover
        raise SystemExit("Missing dependency: pyyaml. Please install the project dependencies first.") from exc

    try:
        from PIL import Image
    except ImportError as exc:  # pragma: no cover
        raise SystemExit("Missing dependency: pillow. Please install the project dependencies first.") from exc

    try:
        import pyarrow.parquet as pq
    except ImportError as exc:  # pragma: no cover
        raise SystemExit("Missing dependency: pyarrow. Please install the project dependencies first.") from exc

    return yaml, Image, pq


def http_get_json(url: str) -> dict[str, Any]:
    with urllib.request.urlopen(url) as response:
        charset = response.headers.get_content_charset() or "utf-8"
        payload = response.read().decode(charset)
    return json.loads(payload)


def fetch_parquet_manifest(repo_id: str, split: str) -> list[dict[str, Any]]:
    url = f"{PARQUET_ENDPOINT}?dataset={urllib.parse.quote(repo_id, safe='')}"
    payload = http_get_json(url)
    if payload.get("failed"):
        raise RuntimeError(f"Parquet manifest request failed: {payload['failed']}")
    files = [item for item in payload.get("parquet_files", []) if item.get("split") == split]
    if not files:
        raise RuntimeError(f"No parquet files found for dataset={repo_id!r}, split={split!r}")
    return files


def download_file(url: str, destination: Path, label: str) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists():
        size_mib = destination.stat().st_size / (1024 ** 2)
        log(f"Reusing existing download: {label} ({size_mib:.1f} MiB)")
        return

    tmp_path = destination.with_name(f"{destination.name}.part")
    resume_bytes = tmp_path.stat().st_size if tmp_path.exists() else 0
    request = urllib.request.Request(url)
    if resume_bytes > 0:
        request.add_header("Range", f"bytes={resume_bytes}-")
        log(f"Resuming download: {label} from {resume_bytes / (1024 ** 2):.1f} MiB")

    try:
        response = urllib.request.urlopen(request)
    except urllib.error.HTTPError as exc:
        if exc.code == 416 and tmp_path.exists():
            os.replace(tmp_path, destination)
            size_mib = destination.stat().st_size / (1024 ** 2)
            log(f"Recovered completed download from partial file: {label} ({size_mib:.1f} MiB)")
            return
        raise

    status = getattr(response, "status", None)
    if resume_bytes > 0 and status != 206:
        log(f"Server did not support resume for {label}; restarting download from scratch")
        tmp_path.unlink()
        resume_bytes = 0
        response.close()
        response = urllib.request.urlopen(url)

    total_bytes = response.headers.get("Content-Length")
    total_bytes = int(total_bytes) if total_bytes is not None else None
    if total_bytes is not None and resume_bytes > 0:
        total_bytes += resume_bytes
    downloaded = resume_bytes
    next_report = max(64 * 1024 * 1024, ((downloaded // (64 * 1024 * 1024)) + 1) * (64 * 1024 * 1024))
    mode = "ab" if resume_bytes > 0 else "wb"
    with response, tmp_path.open(mode) as f:
        if resume_bytes > 0:
            log(f"Continuing download: {label} -> {destination}")
        else:
            log(f"Starting download: {label} -> {destination}")
        while True:
            chunk = response.read(DOWNLOAD_CHUNK_SIZE)
            if not chunk:
                break
            f.write(chunk)
            downloaded += len(chunk)
            if total_bytes is not None and downloaded >= next_report:
                pct = 100.0 * downloaded / total_bytes
                log(f"Downloading {label}: {downloaded / (1024 ** 2):.1f} MiB / {total_bytes / (1024 ** 2):.1f} MiB ({pct:.1f}%)")
                next_report += 64 * 1024 * 1024
    os.replace(tmp_path, destination)
    size_mib = destination.stat().st_size / (1024 ** 2)
    log(f"Finished download: {label} ({size_mib:.1f} MiB)")


def pick_prompt(row: Mapping[str, Any], prompt_column: str) -> str:
    value = row.get(prompt_column)
    if isinstance(value, str) and value.strip():
        return value.strip()
    for key in ("edit", "prompt", "instruction", "text"):
        value = row.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    raise KeyError(f"Could not find a prompt field in row with keys: {sorted(row.keys())}")


def pick_source_image(row: Mapping[str, Any], image_column: str) -> Any:
    value = row.get(image_column)
    if value is not None:
        return value
    for key in ("input_image", "image", "source_image", "source_img"):
        if key in row and row[key] is not None:
            return row[key]
    raise KeyError(f"Could not find a source image field in row with keys: {sorted(row.keys())}")


def read_rows_from_parquet(parquet_path: Path, pq_module: Any, columns: list[str]) -> list[dict[str, Any]]:
    table = pq_module.read_table(parquet_path, columns=columns)
    return table.to_pylist()


def filter_rows(rows: list[dict[str, Any]], args: argparse.Namespace) -> list[dict[str, Any]]:
    if args.dataset_preset == "magicbrush" and args.magicbrush_first_turn_only:
        rows = [row for row in rows if row.get("turn_index") == 1]
    return rows


def load_image_object(image_obj: Any, image_cls: type) -> Any:
    if isinstance(image_obj, image_cls.Image):
        return image_obj
    if isinstance(image_obj, str):
        return image_cls.open(image_obj)
    if isinstance(image_obj, Mapping):
        image_src = image_obj.get("src")
        if image_src:
            with urllib.request.urlopen(image_src) as response:
                return image_cls.open(io.BytesIO(response.read()))
        path = image_obj.get("path")
        if path:
            return image_cls.open(path)
        image_bytes = image_obj.get("bytes")
        if image_bytes is not None:
            return image_cls.open(io.BytesIO(image_bytes))
    raise TypeError(f"Unsupported image object type: {type(image_obj)!r}")


def export_rows(
    rows: list[dict[str, Any]],
    exported: dict[str, dict[str, str]],
    state: dict[str, Any],
    state_path: Path,
    shard_name: str,
    rows_consumed: int,
    image_dir: Path,
    yaml_dir: Path,
    image_cls: type,
    sample_size: int,
    overwrite: bool,
    reuse_existing_images: bool,
    resume_mode: bool,
    image_column: str,
    prompt_column: str,
) -> bool:
    key_width = max(4, len(str(max(sample_size - 1, 0))))
    shard_state = state["shards"].setdefault(shard_name, {})
    shard_state["filtered_row_count"] = len(rows)
    for row_index in range(rows_consumed, len(rows)):
        if len(exported) >= sample_size:
            return True
        row = rows[row_index]
        prompt = pick_prompt(row, prompt_column=prompt_column)
        index = len(exported)
        stem = f"{index:0{key_width}d}"
        image_path = image_dir / f"{stem}.png"
        if image_path.exists():
            if overwrite:
                source_image = pick_source_image(row, image_column=image_column)
                image = load_image_object(source_image, image_cls)
                image.save(image_path)
                image.close()
                log(f"Overwrote image {len(exported) + 1}/{sample_size}: {image_path}")
            elif reuse_existing_images:
                log(f"Reused existing image {len(exported) + 1}/{sample_size}: {image_path}")
            elif resume_mode:
                image_path.unlink()
                log(f"Removed stale uncheckpointed image: {image_path}")
                source_image = pick_source_image(row, image_column=image_column)
                image = load_image_object(source_image, image_cls)
                image.save(image_path)
                image.close()
                log(f"Re-saved image {len(exported) + 1}/{sample_size}: {image_path}")
            else:
                raise SystemExit(
                    f"Image already exists: {image_path}. Pass --overwrite to replace it or "
                    "--reuse-existing-images to keep it."
                )
        else:
            source_image = pick_source_image(row, image_column=image_column)
            image = load_image_object(source_image, image_cls)
            image.save(image_path)
            image.close()
            log(f"Saved image {len(exported) + 1}/{sample_size}: {image_path}")

        rel_path = Path(os.path.relpath(image_path, start=yaml_dir))
        exported[stem] = {"image": rel_path.as_posix(), "prompt": prompt}
        shard_state["rows_consumed"] = row_index + 1
        save_state(state_path, state)
    return len(exported) >= sample_size


def main() -> int:
    args = parse_args()
    if args.sample_size <= 0:
        raise SystemExit("--sample-size must be > 0 for this bounded exporter.")
    if args.overwrite and args.reuse_existing_images:
        raise SystemExit("--overwrite and --reuse-existing-images are mutually exclusive.")

    output_yaml = args.output_yaml.resolve()
    image_dir = args.image_dir.resolve()
    yaml_dir = output_yaml.parent
    preset = resolve_preset(args)
    job_config = build_job_config(args, preset, output_yaml=output_yaml, image_dir=image_dir)
    resume_dir = resolve_resume_dir(job_config)
    state_path = resume_dir / "state.json"
    shard_cache_dir = resolve_shard_cache_dir(preset, args.cache_dir)
    resume_mode = state_path.exists()

    if output_yaml.exists() and not args.overwrite:
        raise SystemExit(f"Output YAML already exists: {output_yaml}. Pass --overwrite to replace it.")
    if image_dir.exists() and any(image_dir.iterdir()) and not (args.overwrite or args.reuse_existing_images or resume_mode):
        raise SystemExit(
            f"Image directory is not empty: {image_dir}. Pass --overwrite to replace files or "
            "--reuse-existing-images to keep them."
        )

    yaml_mod, image_mod, pq_mod = import_dependencies()
    yaml_dir.mkdir(parents=True, exist_ok=True)
    image_dir.mkdir(parents=True, exist_ok=True)
    shard_cache_dir.mkdir(parents=True, exist_ok=True)
    state = load_state(state_path, job_config)
    exported = state["exported"]

    log("Phase 1/4: fetching parquet shard manifest")
    log(f"Dataset preset: {args.dataset_preset}")
    log(f"Dataset: {preset.repo_id} [{preset.default_split}]")
    log(f"Resume workspace: {resume_dir}")
    log(f"Shard cache: {shard_cache_dir}")
    manifest = fetch_parquet_manifest(preset.repo_id, preset.default_split)
    log(f"Found {len(manifest)} parquet shards for split {preset.default_split}")

    log("Phase 2/4: preparing shard order")
    rng = random.Random(args.seed)
    if args.shuffle:
        log(f"Shuffling shard order with seed={args.seed}")
        rng.shuffle(manifest)
    else:
        log("Using shard order from the manifest")
    if args.dataset_preset == "magicbrush" and args.magicbrush_first_turn_only:
        log("MagicBrush filter enabled: only keeping rows with turn_index == 1")
    if exported:
        log(f"Resuming from checkpoint with {len(exported)}/{args.sample_size} samples already exported")

    log("Phase 3/4: downloading shards and exporting images")
    for shard_index, shard in enumerate(manifest, start=1):
        if len(exported) >= args.sample_size:
            break
        shard_name = shard["filename"]
        shard_path = shard_cache_dir / shard_name
        shard_state = state["shards"].setdefault(shard_name, {})
        rows_consumed = int(shard_state.get("rows_consumed", 0))
        filtered_row_count = shard_state.get("filtered_row_count")
        if rows_consumed and filtered_row_count is not None and rows_consumed >= filtered_row_count:
            log(f"Skipping completed shard {shard_index}/{len(manifest)}: {shard_name}")
            if args.shuffle:
                advance_rng_for_shuffle(rng, int(filtered_row_count))
            continue
        if shard_path.exists():
            log(f"Reusing cached shard {shard_index}/{len(manifest)}: {shard_path}")
        else:
            size_mib = shard.get("size", 0) / (1024 ** 2)
            log(f"Preparing shard {shard_index}/{len(manifest)}: {shard_name} ({size_mib:.1f} MiB)")
            download_file(shard["url"], shard_path, shard_name)

        columns = [preset.image_column, preset.prompt_column, *preset.extra_columns]
        for read_attempt in range(2):
            log(f"Reading rows from shard {shard_name}")
            try:
                rows = read_rows_from_parquet(shard_path, pq_mod, columns=columns)
                break
            except Exception as exc:
                if read_attempt == 1:
                    raise
                log(f"Failed to read shard {shard_name}: {exc}")
                if shard_path.exists():
                    shard_path.unlink()
                    log(f"Removed corrupted shard cache: {shard_path}")
                size_mib = shard.get("size", 0) / (1024 ** 2)
                log(f"Re-downloading shard {shard_index}/{len(manifest)}: {shard_name} ({size_mib:.1f} MiB)")
                download_file(shard["url"], shard_path, shard_name)
        log(f"Loaded {len(rows)} rows from shard {shard_name}")
        rows = filter_rows(rows, args)
        log(f"Kept {len(rows)} rows after dataset-specific filtering")
        shard_state["filtered_row_count"] = len(rows)
        save_state(state_path, state)
        if args.shuffle:
            rng.shuffle(rows)
            log(f"Shuffled rows inside shard {shard_name}")

        done = export_rows(
            rows=rows,
            exported=exported,
            state=state,
            state_path=state_path,
            shard_name=shard_name,
            rows_consumed=rows_consumed,
            image_dir=image_dir,
            yaml_dir=yaml_dir,
            image_cls=image_mod,
            sample_size=args.sample_size,
            overwrite=args.overwrite,
            reuse_existing_images=args.reuse_existing_images,
            resume_mode=resume_mode,
            image_column=preset.image_column,
            prompt_column=preset.prompt_column,
        )
        if done:
            break

    exported_count = len(exported)
    if exported_count == 0:
        raise RuntimeError("Export completed without any samples; refusing to write an empty YAML file.")
    if exported_count < args.sample_size:
        log(
            f"Warning: only exported {exported_count} samples, fewer than requested {args.sample_size}. "
            "Writing partial YAML with the available samples."
        )

    log("Phase 4/4: writing YAML file")
    with output_yaml.open("w", encoding="utf-8") as f:
        yaml_mod.safe_dump(exported, f, sort_keys=False, allow_unicode=True, width=120)

    cleanup_temp_artifacts(resume_dir)
    log(f"Removed resume workspace: {resume_dir}")
    if args.keep_shards:
        log(f"Kept shard cache: {shard_cache_dir}")

    log(f"Wrote YAML: {output_yaml}")
    log(f"Saved images directory: {image_dir}")
    log(f"Completed export with {exported_count} samples")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Thanks for sharing the data! I'll dig into the details.

I've actually finished quantizing the Klein model recently. One interesting thing I noticed is that even calibrating with only T2I data still preserves its editing capability quite well. But I do think mixing in some edit data would probably make the calibration results even better.