code/models/tt_dit/tests/dataset_eval/reference/motif.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /tests /dataset_eval /reference /motif.py

Add files using upload-large-folder tool

33d23e1 verified 12 days ago

6.9 kB

# 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()

	# 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()