code/models/tt_dit/tests/dataset_eval/reference/motif.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

import argparse
import csv
import json
import os
import subprocess
import time
import urllib
from pathlib import Path

from loguru import logger
from PIL import Image

from ..clip_encoder import CLIPEncoder
from ..fid_score import calculate_fid_score

COCO_CAPTIONS_DOWNLOAD_PATH = (
    "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/"
    "text_to_image/coco2014/captions/captions_source.tsv"
)


def load_coco_prompts(captions_path, start_from=0, num_prompts=16):
    """loads prompts from COCO captions tsv, downloading if missing"""
    prompts = []
    if not os.path.isfile(captions_path):
        logger.info(f"file {captions_path} not found. downloading...")
        urllib.request.urlretrieve(COCO_CAPTIONS_DOWNLOAD_PATH, captions_path)
        logger.info("download complete.")

    with open(captions_path, "r") as tsv_file:
        reader = csv.reader(tsv_file, delimiter="\t")
        next(reader)  # skip header
        for index, row in enumerate(reader):
            if index < start_from:
                continue
            if index >= start_from + num_prompts:
                break
            if len(row) >= 3:
                prompts.append(row[2].strip())
    logger.info(f"loaded {len(prompts)} coco prompts from {captions_path}")
    return prompts


def generate_images_with_inference(prompts_file, output_dir, num_prompts):
    """forces motif inference to regenerate only the requested number of prompts"""
    logger.info("clearing any previous images...")
    os.makedirs(output_dir, exist_ok=True)
    for f in Path(output_dir).rglob("*.png"):
        f.unlink()

    # create temporary subset prompt file
    tmp_prompts_file = os.path.join(output_dir, "subset_prompts.txt")
    with open(prompts_file, "r") as f_in, open(tmp_prompts_file, "w") as f_out:
        lines = [line.strip() for line in f_in if line.strip()]
        subset = lines[:num_prompts]
        for l in subset:
            f_out.write(l + "\n")

    logger.info(f"generating {len(subset)} images using motif inference.py ...")

    cmd = [
        "python",
        "inference.py",
        "--model-config",
        "configs/mmdit_xlarge_hq.json",
        "--model-ckpt",
        "checkpoints/pytorch_model_fsdp.bin",
        "--seed",
        "7777",
        "--steps",
        "20",
        "--resolution",
        "1024",
        "--prompt-file",
        tmp_prompts_file,  # use subset
        "--guidance-scales",
        "4.0",
        "--output-dir",
        output_dir,
        "--batch-size",
        "1",
    ]
    subprocess.run(cmd, check=True)
    logger.info("image generation complete.")


def compute_clip_scores(images_dir, prompts):
    """computes average clip similarity across generated images and corresponding prompts"""
    clip_encoder = CLIPEncoder(clip_version="ViT-B/32", pretrained="openai")
    scores = []

    # find all pngs recursively
    all_images = sorted(Path(images_dir).rglob("*.png"))
    if len(all_images) < len(prompts):
        logger.warning(f"found only {len(all_images)} images, expected {len(prompts)}")

    for idx, (prompt, img_path) in enumerate(zip(prompts, all_images)):
        try:
            image = Image.open(img_path).convert("RGB")
            score = clip_encoder.get_clip_score(prompt, image)
            scores.append(score.item())
        except Exception as e:
            logger.error(f"error processing {img_path}: {e}")

    if not scores:
        return 0.0

    avg_score = sum(scores) / len(scores)
    logger.info(f"average clip score = {avg_score:.4f} ({avg_score * 100:.2f} CLIPScore)")
    return avg_score


def compute_fid(images_dir, coco_stats="val2014.npz"):
    """computes fid score against coco statistics"""
    if not os.path.exists(coco_stats):
        logger.info(f"COCO stats not found at {coco_stats}, downloading...")
        os.makedirs(os.path.dirname(coco_stats) or ".", exist_ok=True)
        urllib.request.urlretrieve(
            "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/"
            "text_to_image/tools/val2014.npz",
            coco_stats,
        )
        logger.info("download complete.")

    all_images = [Image.open(f).convert("RGB") for f in sorted(Path(images_dir).rglob("*.png"))]
    if not all_images:
        logger.error("no images found for fid computation")
        return None

    logger.info(f"computing fid over {len(all_images)} images ...")
    fid = calculate_fid_score(all_images, coco_stats)
    logger.info(f"fid score = {fid:.4f}")
    return fid


def main():
    parser = argparse.ArgumentParser(description="Motif Dataset Evaluation")
    parser.add_argument("--images-dir", type=str, required=True, help="directory containing generated images")
    parser.add_argument("--prompts-file", type=str, required=True, help="path to prompts file (.tsv or .txt)")
    parser.add_argument("--output-json", type=str, default="dataset_eval_results.json", help="output results json")
    parser.add_argument("--start-from", type=int, default=0, help="starting index for captions.tsv")
    parser.add_argument("--num-prompts", type=int, default=16, help="number of prompts to evaluate")
    args = parser.parse_args()

    start_time = time.time()

    # load prompts dynamically depending on file type
    if args.prompts_file.endswith(".tsv"):
        prompts = load_coco_prompts(args.prompts_file, args.start_from, args.num_prompts)
    else:
        with open(args.prompts_file, "r") as f:
            prompts = [line.strip() for line in f if line.strip()]

    # always regenerate fresh images
    generate_images_with_inference(args.prompts_file, args.images_dir, args.num_prompts)

    # automatically detect where the actual images were saved
    def find_deepest_image_dir(root):
        """walks the output dir and returns the first subdir containing pngs"""
        for dirpath, _, filenames in os.walk(root):
            if any(fn.lower().endswith(".png") for fn in filenames):
                return dirpath
        return root

    image_root = find_deepest_image_dir(args.images_dir)
    logger.info(f"found generated images under: {image_root}")

    # compute metrics at the detected path
    clip_score = compute_clip_scores(image_root, prompts)
    fid_score = compute_fid(image_root, coco_stats="val2014.npz")
    total_time = time.time() - start_time

    results = {
        "num_images": len(prompts),
        "clip_score": clip_score,
        "fid_score": fid_score,
        "total_time_sec": total_time,
        "images_dir": image_root,
    }

    with open(args.output_json, "w") as f:
        json.dump(results, f, indent=2)

    logger.info(f"saved results to {args.output_json}")
    logger.info(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()