code/models/tt_dit/tests/dataset_eval/test_tt_dit.py

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

import csv
import itertools
import json
import os
import random
import statistics
import time
import urllib

import pytest
from loguru import logger

import ttnn

# from .fid_score import calculate_fid_score
from models.demos.stable_diffusion_xl_base.utils.fid_score import calculate_fid_score
from models.perf.benchmarking_utils import BenchmarkProfiler
from models.tt_dit.pipelines.events import profiler_event_callback

from .clip_encoder import CLIPEncoder

# COCO captions download path
COCO_CAPTIONS_DOWNLOAD_PATH = "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv"


def get_prompts(captions_path, start_from, num_prompts):
    assert (
        0 <= start_from < 5000 and start_from + num_prompts <= 5000
    ), "start_from must be between 0 and 4999, and start_from + num_prompts must not exceed 5000."

    prompts = []
    if not os.path.isfile(captions_path):
        logger.info(f"file {captions_path} not found. downloading...")
        os.makedirs(os.path.dirname(captions_path), exist_ok=True)
        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
            prompts.append(row[2])
    return prompts


def get_img_evals(model_id, model_setup, images, prompts, coco_statistics_path):
    # Calculate metrics
    logger.info("Calculating CLIP scores...")
    num_prompts = len(prompts)
    clip = CLIPEncoder()
    clip_scores = [100 * clip.get_clip_score(prompts[i], img).item() for i, img in enumerate(images)]
    average_clip_score = sum(clip_scores) / len(clip_scores)

    deviation_clip_score = "N/A"
    fid_score = "N/A"

    # Calculate CLIP standard deviation if we have multiple prompts
    if num_prompts >= 2:
        deviation_clip_score = statistics.stdev(clip_scores)
        logger.info(f"CLIP scores individual: {[f'{score:.2f}' for score in clip_scores]}")

        # Calculate FID score if we have the COCO statistics file
        logger.info("Calculating FID score...")
        try:
            fid_score = calculate_fid_score(images, coco_statistics_path)
        except Exception as e:
            logger.error(f"Error calculating FID score: {e}")
    else:
        logger.info("FID calculation requires at least 2 images")

    evals = [
        {
            "model": model_id,
            "average_clip": average_clip_score,
            "clip_standard_deviation": deviation_clip_score,
            "num_prompts": num_prompts,
            "accuracy_check": check_accuracy(num_prompts, average_clip_score, "clip", model_setup),
        },
        {
            "model": model_id,
            "fid_score": fid_score,
            "num_prompts": num_prompts,
            "accuracy_check": check_accuracy(num_prompts, fid_score, "fid", model_setup),
        },
    ]
    return evals


def get_vid_evals(model_id, model_setup, videos, prompts):
    num_prompts = len(prompts)
    clip = CLIPEncoder()
    clip_scores = [
        [100 * clip.get_clip_score(prompts[i], img).item() for img in video] for i, video in enumerate(videos)
    ]
    clip_score_stats = [(min(c), max(c), statistics.mean(c), statistics.stdev(c)) for c in clip_scores]

    for p, c in zip(prompts, clip_score_stats):
        logger.info(f"Prompt: {p}, CLIP statistics (min, max, mean, stddev): {c}")

    video_mean_scores = [c[2] for c in clip_score_stats]
    min_clip = min(video_mean_scores)
    max_clip = max(video_mean_scores)
    mean_clip = statistics.mean(video_mean_scores)
    stddev_clip = statistics.stdev(video_mean_scores)

    evals = [
        {
            "model": model_id,
            "average_clip": mean_clip,
            "min_clip": min_clip,
            "max_clip": max_clip,
            "clip_standard_deviation": stddev_clip,
            "num_prompts": num_prompts,
            "accuracy_check": check_accuracy(num_prompts, mean_clip, "clip", model_setup),
        }
    ]
    return evals


def check_accuracy(num_prompts, score, score_key, model_info):
    accuracy_status = 1
    prompt_data = model_info["accuracy"].get(str(num_prompts), {})
    range_key = f"{score_key}_valid_range"
    if range_key in prompt_data:
        score_range = prompt_data[range_key]
        if score >= score_range[0] and score <= score_range[1]:
            return 2
        else:
            return 3

    return accuracy_status


@pytest.mark.parametrize(
    "mesh_device,device_name",
    [
        [(1, 2), "p300"],
        [(2, 2), "qbge"],
        [(2, 4), "t3K"],
        [(4, 8), "galaxy"],
    ],
    ids=[
        "1x2",
        "2x2",
        "2x4",
        "4x8",
    ],
    indirect=["mesh_device"],
)
def test_tt_dit_accuracy(
    *,
    mesh_device: ttnn.MeshDevice,
    device_name: str,
    num_inference_steps,
    model_id,
    model_location_generator,
    dit_pipeline,
    model_metadata,
    evaluation_range,
    is_video,
    is_24hrs_test,
    monkeypatch,
) -> None:
    """Accuracy test for TT-Metal DiT pipelines with CLIP and FID score evaluation.
    It relies on the default configurations set in each pipeline.
    """

    if os.getenv("TT_DIT_CACHE_DIR") is None:
        monkeypatch.setenv("TT_DIT_CACHE_DIR", "/tmp/TT_DIT_CACHE")
    benchmark_profiler = BenchmarkProfiler()
    if ttnn.device.is_blackhole() and device_name == "galaxy":
        device_name = "bh_" + device_name
    captions_path = "coco_captions/captions.tsv"
    coco_statistics_path = "coco_statistics/val2014.npz"
    start_from, num_prompts = evaluation_range
    prompts = get_prompts(captions_path, start_from, num_prompts)
    logger.info(f"start inference from prompt index: {start_from} to {start_from + num_prompts}")

    logger.info(f"  Mesh shape: {mesh_device.shape}")
    logger.info(f"  Model ID: {model_id}")
    logger.info(f"  Inference steps: {num_inference_steps}")

    # Create pipeline similar to test_performance_flux1.py
    pipeline = dit_pipeline.create_pipeline(
        mesh_device=mesh_device, checkpoint_name=model_location_generator(model_metadata["hf_id"])
    )

    images = []
    gen_type = "image" if not is_video else "video"
    logger.info(f"Starting generation of {len(prompts)} {gen_type}s...")

    # Generate images for each prompt
    iter = itertools.cycle(prompts) if is_24hrs_test else prompts
    start_time = time.time()
    for i, prompt in enumerate(iter):
        logger.info(
            f"Generating {gen_type} {i+1} of {len(prompts) if not is_24hrs_test else '24hrs {gen_type}s'}: {prompt[:50]}..."
        )

        # Generate image using TT-Metal pipeline
        with benchmark_profiler("inference", iteration=i):
            generated_images = pipeline.run_single_prompt(
                prompt=prompt,
                num_inference_steps=num_inference_steps,
                on_event=profiler_event_callback(benchmark_profiler, i),
                **model_metadata.get("extra_args", {}),
            )

        images.append(generated_images[0])
        logger.info(f"{gen_type} {i+1} completed in {benchmark_profiler.get_duration('inference', i):.2f}s")

        if is_24hrs_test:
            if (time.time() - start_time) > 24 * 60 * 60:
                break
            # Sleep for a random time between 0.1 and 10 seconds
            time.sleep(random.uniform(0.1, 10))

    # Performance metrics
    average_inference_time = benchmark_profiler.get_duration_average("inference", 1)  # Skip warmup iterations

    data = {
        "model": model_id,
        "metadata": {
            "device": f"{device_name}",
            "model_name": model_id,
            "start_from": start_from,
            "num_prompts": num_prompts,
            "num_inference_steps": num_inference_steps,
            "mesh_shape": list(mesh_device.shape),
        },
        "benchmarks_summary": [
            {
                "device": f"{device_name}",
                "model": model_id,
                "average_inference_time": average_inference_time,
                "average_vae_time": benchmark_profiler.get_duration_average("vae", 1),
                "average_encoding_time": benchmark_profiler.get_duration_average("encoder", 1),
                "average_denoising_time": benchmark_profiler.get_duration_average("denoising", 1),
            }
        ],
        "evals": get_vid_evals(model_id, model_metadata, images, prompts)
        if is_video
        else get_img_evals(model_id, model_metadata, images, prompts, coco_statistics_path),
    }

    out_root = "test_reports"
    os.makedirs(out_root, exist_ok=True)
    file_path = f"{out_root}/sdxl_test_results.json"
    with open(file_path, "w") as f:
        json.dump(data, f, indent=4)
    logger.info(f"test results saved to {file_path}")

    # Synchronize devices
    pipeline.synchronize_devices()

    print(f"\n🎉 Test completed successfully!")
    print(f"⚡ Average Inference Time: {average_inference_time:.2f}s")