| |
| |
|
|
| import pytest |
|
|
| import ttnn |
|
|
| from ...pipelines.flux1.pipeline_flux1 import Flux1Pipeline |
| from ...pipelines.mochi.pipeline_mochi import MochiPipeline |
| from ...pipelines.motif.pipeline_motif import MotifPipeline |
| from ...pipelines.stable_diffusion_35_large.pipeline_stable_diffusion_35_large import StableDiffusion3Pipeline |
| from ...pipelines.wan.pipeline_wan import WanPipeline |
|
|
| targets_setup = { |
| "stable-diffusion-3.5-large": { |
| "hf_id": "stabilityai/stable-diffusion-3.5-large", |
| "num_inference_steps": 28, |
| "pipeline": StableDiffusion3Pipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| "l1_small_size": 32768, |
| "trace_region_size": 25000000, |
| }, |
| "accuracy": { |
| "5000": {"fid_valid_range": [22.63592, 24.0361780112], "clip_valid_range": [30.38525, 32.264881565]}, |
| "500": {"fid_valid_range": [79.734, 84.66634524], "clip_valid_range": [30.423468, 32.30546373048]}, |
| "100": { |
| "fid_valid_range": [175.35951, 186.2072492886], |
| "clip_valid_range": [30.162635, 32.028495601100005], |
| }, |
| }, |
| }, |
| "flux.1-dev": { |
| "hf_id": "black-forest-labs/FLUX.1-dev", |
| "num_inference_steps": 28, |
| "pipeline": Flux1Pipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| "l1_small_size": 32768, |
| "trace_region_size": 50000000, |
| }, |
| "accuracy": { |
| "5000": {"fid_valid_range": [28.14067, 29.8814518462], "clip_valid_range": [29.7014, 31.538728604]}, |
| "500": {"fid_valid_range": [82.37143, 87.46692665980001], "clip_valid_range": [29.79258, 31.6355489988]}, |
| "100": {"fid_valid_range": [174.97054, 185.79421760440002], "clip_valid_range": [29.86339, 31.7107393054]}, |
| }, |
| }, |
| "flux.1-schnell": { |
| "hf_id": "black-forest-labs/FLUX.1-schnell", |
| "num_inference_steps": 4, |
| "pipeline": Flux1Pipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| "l1_small_size": 32768, |
| "trace_region_size": 50000000, |
| }, |
| "accuracy": { |
| "5000": { |
| "fid_valid_range": [22.195637, 23.568659104820004], |
| "clip_valid_range": [30.010539, 31.866990942540003], |
| }, |
| "500": {"fid_valid_range": [77.21976, 81.9965743536], "clip_valid_range": [30.19901, 32.067120758600005]}, |
| "100": {"fid_valid_range": [171.065611, 181.64772969646], "clip_valid_range": [30.41629, 32.2978416994]}, |
| }, |
| }, |
| "motif-image-6b-preview": { |
| "hf_id": "Motif-Technologies/Motif-Image-6B-Preview", |
| "num_inference_steps": 28, |
| "pipeline": MotifPipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| "l1_small_size": 32768, |
| "trace_region_size": 31000000, |
| }, |
| "accuracy": { |
| "5000": {"fid_valid_range": [33.85203, 35.9461165758], "clip_valid_range": [29.6529, 31.487228394]}, |
| "500": {"fid_valid_range": [86.02833, 91.3500424938], "clip_valid_range": [29.65581, 31.4903184066]}, |
| "100": {"fid_valid_range": [180.97872, 192.1740636192], "clip_valid_range": [29.65581, 31.4903184066]}, |
| }, |
| }, |
| "mochi-1-preview": { |
| "hf_id": "genmo/mochi-1-preview", |
| "num_inference_steps": 50, |
| "pipeline": MochiPipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| }, |
| "accuracy": {"16": {"clip_valid_range": [26.20977, 27.83099]}}, |
| "is_video": True, |
| "extra_args": {"num_frames": 168}, |
| }, |
| "wan2.2": { |
| "hf_id": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", |
| "num_inference_steps": 40, |
| "pipeline": WanPipeline, |
| "device_params": { |
| "fabric_config": ttnn.FabricConfig.FABRIC_1D, |
| }, |
| "accuracy": {"16": {"clip_valid_range": [28.25276, 30.00036]}}, |
| "is_video": True, |
| "extra_args": {"output_type": "pil"}, |
| }, |
| } |
|
|
|
|
| def pytest_addoption(parser): |
| parser.addoption( |
| "--start-from", |
| action="store", |
| default=0, |
| help="Start from prompt number (0-4999)", |
| ) |
| parser.addoption( |
| "--num-prompts", |
| action="store", |
| default=5000, |
| help="Number of prompts to process (default: 5000)", |
| ) |
|
|
| model_id_choices = targets_setup.keys() |
| parser.addoption( |
| "--model-id", |
| action="store", |
| required=True, |
| choices=model_id_choices, |
| help=f"Model ID to use for evaluation. Options:{','.join(model_id_choices)}", |
| ) |
|
|
| parser.addoption( |
| "--num-inference-steps", |
| action="store", |
| default=None, |
| type=int, |
| help="Number of inference steps (default: 28)", |
| ) |
|
|
| parser.addoption( |
| "--24hrs-test", |
| action="store_true", |
| default=False, |
| help="Whether to run a 24-hour test", |
| ) |
|
|
|
|
| @pytest.fixture |
| def evaluation_range(request): |
| start_from = request.config.getoption("--start-from") |
| num_prompts = request.config.getoption("--num-prompts") |
| if start_from is not None: |
| start_from = int(start_from) |
| else: |
| start_from = 0 |
| if num_prompts is not None: |
| num_prompts = int(num_prompts) |
| else: |
| num_prompts = 5 |
| return start_from, num_prompts |
|
|
|
|
| @pytest.fixture |
| def model_id(request): |
| return request.config.getoption("--model-id") |
|
|
|
|
| @pytest.fixture |
| def is_24hrs_test(request): |
| return request.config.getoption("--24hrs-test") |
|
|
|
|
| @pytest.fixture |
| def num_inference_steps(request): |
| return ( |
| request.config.getoption("--num-inference-steps") |
| or targets_setup[request.config.getoption("--model-id")]["num_inference_steps"] |
| ) |
|
|
|
|
| @pytest.fixture |
| def dit_pipeline(model_id): |
| return targets_setup[model_id]["pipeline"] |
|
|
|
|
| @pytest.fixture |
| def model_metadata(model_id): |
| return targets_setup[model_id] |
|
|
|
|
| @pytest.fixture |
| def device_params(request, model_id): |
| """Return device_params based on model_id, with optional parametrization override.""" |
|
|
| |
| params = getattr(request, "param", {}).copy() |
|
|
| |
| if not params: |
| params = targets_setup[model_id]["device_params"] |
| |
| else: |
| model_defaults = targets_setup[model_id]["device_params"] |
| params = {**model_defaults, **params} |
|
|
| return params |
|
|
|
|
| @pytest.fixture |
| def is_video(model_id): |
| return targets_setup[model_id].get("is_video", False) |
|
|
