cache-dit 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of cache-dit might be problematic. Click here for more details.
- cache_dit/_version.py +2 -2
- cache_dit/cache_factory/dual_block_cache/cache_context.py +19 -17
- cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py +1 -1
- cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py +1 -1
- cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py +1 -1
- cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py +1 -1
- cache_dit/metrics/__init__.py +14 -0
- cache_dit/metrics/config.py +34 -0
- cache_dit/metrics/fid.py +498 -0
- cache_dit/metrics/inception.py +353 -0
- cache_dit/metrics/metrics.py +469 -0
- {cache_dit-0.2.5.dist-info → cache_dit-0.2.7.dist-info}/METADATA +37 -1
- {cache_dit-0.2.5.dist-info → cache_dit-0.2.7.dist-info}/RECORD +17 -11
- cache_dit-0.2.7.dist-info/entry_points.txt +2 -0
- {cache_dit-0.2.5.dist-info → cache_dit-0.2.7.dist-info}/WHEEL +0 -0
- {cache_dit-0.2.5.dist-info → cache_dit-0.2.7.dist-info}/licenses/LICENSE +0 -0
- {cache_dit-0.2.5.dist-info → cache_dit-0.2.7.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,469 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import cv2
|
|
3
|
+
import pathlib
|
|
4
|
+
import argparse
|
|
5
|
+
import numpy as np
|
|
6
|
+
from tqdm import tqdm
|
|
7
|
+
from functools import partial
|
|
8
|
+
from skimage.metrics import mean_squared_error
|
|
9
|
+
from skimage.metrics import peak_signal_noise_ratio
|
|
10
|
+
from skimage.metrics import structural_similarity
|
|
11
|
+
from cache_dit.metrics.fid import FrechetInceptionDistance
|
|
12
|
+
from cache_dit.metrics.config import set_metrics_verbose
|
|
13
|
+
from cache_dit.metrics.config import get_metrics_verbose
|
|
14
|
+
from cache_dit.metrics.config import _IMAGE_EXTENSIONS
|
|
15
|
+
from cache_dit.metrics.config import _VIDEO_EXTENSIONS
|
|
16
|
+
from cache_dit.logger import init_logger
|
|
17
|
+
|
|
18
|
+
logger = init_logger(__name__)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
DISABLE_VERBOSE = not get_metrics_verbose()
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def compute_psnr_file(
|
|
25
|
+
image_true: np.ndarray | str,
|
|
26
|
+
image_test: np.ndarray | str,
|
|
27
|
+
) -> float:
|
|
28
|
+
"""
|
|
29
|
+
img_true = cv2.imread(img_true_file)
|
|
30
|
+
img_test = cv2.imread(img_test_file)
|
|
31
|
+
PSNR = compute_psnr(img_true, img_test)
|
|
32
|
+
"""
|
|
33
|
+
if isinstance(image_true, str):
|
|
34
|
+
image_true = cv2.imread(image_true)
|
|
35
|
+
if isinstance(image_test, str):
|
|
36
|
+
image_test = cv2.imread(image_test)
|
|
37
|
+
return peak_signal_noise_ratio(
|
|
38
|
+
image_true,
|
|
39
|
+
image_test,
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def compute_mse_file(
|
|
44
|
+
image_true: np.ndarray | str,
|
|
45
|
+
image_test: np.ndarray | str,
|
|
46
|
+
) -> float:
|
|
47
|
+
"""
|
|
48
|
+
img_true = cv2.imread(img_true_file)
|
|
49
|
+
img_test = cv2.imread(img_test_file)
|
|
50
|
+
MSE = compute_mse(img_true, img_test)
|
|
51
|
+
"""
|
|
52
|
+
if isinstance(image_true, str):
|
|
53
|
+
image_true = cv2.imread(image_true)
|
|
54
|
+
if isinstance(image_test, str):
|
|
55
|
+
image_test = cv2.imread(image_test)
|
|
56
|
+
return mean_squared_error(
|
|
57
|
+
image_true,
|
|
58
|
+
image_test,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def compute_ssim_file(
|
|
63
|
+
image_true: np.ndarray | str,
|
|
64
|
+
image_test: np.ndarray | str,
|
|
65
|
+
) -> float:
|
|
66
|
+
"""
|
|
67
|
+
img_true = cv2.imread(img_true_file)
|
|
68
|
+
img_test = cv2.imread(img_test_file)
|
|
69
|
+
SSIM = compute_ssim(img_true, img_test)
|
|
70
|
+
"""
|
|
71
|
+
if isinstance(image_true, str):
|
|
72
|
+
image_true = cv2.imread(image_true)
|
|
73
|
+
if isinstance(image_test, str):
|
|
74
|
+
image_test = cv2.imread(image_test)
|
|
75
|
+
return structural_similarity(
|
|
76
|
+
image_true,
|
|
77
|
+
image_test,
|
|
78
|
+
multichannel=True,
|
|
79
|
+
channel_axis=2,
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def compute_dir_metric(
|
|
84
|
+
image_true_dir: np.ndarray | str,
|
|
85
|
+
image_test_dir: np.ndarray | str,
|
|
86
|
+
compute_file_func: callable = compute_psnr_file,
|
|
87
|
+
) -> float:
|
|
88
|
+
# Image
|
|
89
|
+
if isinstance(image_true_dir, np.ndarray) or isinstance(
|
|
90
|
+
image_test_dir, np.ndarray
|
|
91
|
+
):
|
|
92
|
+
return compute_file_func(image_true_dir, image_test_dir), 1
|
|
93
|
+
# File
|
|
94
|
+
if not os.path.isdir(image_true_dir) or not os.path.isdir(image_test_dir):
|
|
95
|
+
return compute_file_func(image_true_dir, image_test_dir), 1
|
|
96
|
+
# Dir
|
|
97
|
+
image_true_dir: pathlib.Path = pathlib.Path(image_true_dir)
|
|
98
|
+
image_true_files = sorted(
|
|
99
|
+
[
|
|
100
|
+
file
|
|
101
|
+
for ext in _IMAGE_EXTENSIONS
|
|
102
|
+
for file in image_true_dir.rglob("*.{}".format(ext))
|
|
103
|
+
]
|
|
104
|
+
)
|
|
105
|
+
image_test_dir: pathlib.Path = pathlib.Path(image_test_dir)
|
|
106
|
+
image_test_files = sorted(
|
|
107
|
+
[
|
|
108
|
+
file
|
|
109
|
+
for ext in _IMAGE_EXTENSIONS
|
|
110
|
+
for file in image_test_dir.rglob("*.{}".format(ext))
|
|
111
|
+
]
|
|
112
|
+
)
|
|
113
|
+
image_true_files = [file.as_posix() for file in image_true_files]
|
|
114
|
+
image_test_files = [file.as_posix() for file in image_test_files]
|
|
115
|
+
|
|
116
|
+
# select valid files
|
|
117
|
+
image_true_files_selected = []
|
|
118
|
+
image_test_files_selected = []
|
|
119
|
+
for i in range(min(len(image_true_files), len(image_test_files))):
|
|
120
|
+
selected_image_true = image_true_files[i]
|
|
121
|
+
selected_image_test = image_test_files[i]
|
|
122
|
+
# Image pair must have the same basename
|
|
123
|
+
if os.path.basename(selected_image_test) == os.path.basename(
|
|
124
|
+
selected_image_true
|
|
125
|
+
):
|
|
126
|
+
image_true_files_selected.append(selected_image_true)
|
|
127
|
+
image_test_files_selected.append(selected_image_test)
|
|
128
|
+
image_true_files = image_true_files_selected.copy()
|
|
129
|
+
image_test_files = image_test_files_selected.copy()
|
|
130
|
+
if len(image_true_files) == 0:
|
|
131
|
+
logger.error(
|
|
132
|
+
"No valid Image pairs, please note that Image "
|
|
133
|
+
"pairs must have the same basename."
|
|
134
|
+
)
|
|
135
|
+
return None, None
|
|
136
|
+
|
|
137
|
+
logger.debug(f"image_true_files: {image_true_files}")
|
|
138
|
+
logger.debug(f"image_test_files: {image_test_files}")
|
|
139
|
+
|
|
140
|
+
total_metric = 0.0
|
|
141
|
+
valid_files = 0
|
|
142
|
+
for image_true, image_test in tqdm(
|
|
143
|
+
zip(image_true_files, image_test_files),
|
|
144
|
+
total=len(image_true_files),
|
|
145
|
+
disable=DISABLE_VERBOSE,
|
|
146
|
+
):
|
|
147
|
+
metric = compute_file_func(image_true, image_test)
|
|
148
|
+
if metric != float("inf"):
|
|
149
|
+
total_metric += metric
|
|
150
|
+
valid_files += 1
|
|
151
|
+
|
|
152
|
+
if valid_files > 0:
|
|
153
|
+
average_metric = total_metric / valid_files
|
|
154
|
+
logger.debug(f"Average: {average_metric:.2f}")
|
|
155
|
+
return average_metric, valid_files
|
|
156
|
+
else:
|
|
157
|
+
logger.debug("No valid files to compare")
|
|
158
|
+
return None, None
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def _fetch_video_frames(
|
|
162
|
+
video_true: str,
|
|
163
|
+
video_test: str,
|
|
164
|
+
):
|
|
165
|
+
cap1 = cv2.VideoCapture(video_true)
|
|
166
|
+
cap2 = cv2.VideoCapture(video_test)
|
|
167
|
+
|
|
168
|
+
if not cap1.isOpened() or not cap2.isOpened():
|
|
169
|
+
logger.error("Could not open video files")
|
|
170
|
+
return [], [], 0
|
|
171
|
+
|
|
172
|
+
frame_count = min(
|
|
173
|
+
int(cap1.get(cv2.CAP_PROP_FRAME_COUNT)),
|
|
174
|
+
int(cap2.get(cv2.CAP_PROP_FRAME_COUNT)),
|
|
175
|
+
)
|
|
176
|
+
|
|
177
|
+
valid_frames = 0
|
|
178
|
+
video_true_frames = []
|
|
179
|
+
video_test_frames = []
|
|
180
|
+
|
|
181
|
+
logger.debug(f"Total frames: {frame_count}")
|
|
182
|
+
|
|
183
|
+
while True:
|
|
184
|
+
ret1, frame1 = cap1.read()
|
|
185
|
+
ret2, frame2 = cap2.read()
|
|
186
|
+
|
|
187
|
+
if not ret1 or not ret2:
|
|
188
|
+
break
|
|
189
|
+
|
|
190
|
+
video_true_frames.append(frame1)
|
|
191
|
+
video_test_frames.append(frame2)
|
|
192
|
+
|
|
193
|
+
valid_frames += 1
|
|
194
|
+
|
|
195
|
+
cap1.release()
|
|
196
|
+
cap2.release()
|
|
197
|
+
|
|
198
|
+
if valid_frames <= 0:
|
|
199
|
+
return [], [], 0
|
|
200
|
+
|
|
201
|
+
return video_true_frames, video_test_frames, valid_frames
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def compute_video_metric(
|
|
205
|
+
video_true: str,
|
|
206
|
+
video_test: str,
|
|
207
|
+
compute_frame_func: callable = compute_psnr_file,
|
|
208
|
+
) -> float:
|
|
209
|
+
"""
|
|
210
|
+
video_true = "video_true.mp4"
|
|
211
|
+
video_test = "video_test.mp4"
|
|
212
|
+
PSNR = compute_video_psnr(video_true, video_test)
|
|
213
|
+
"""
|
|
214
|
+
if os.path.isfile(video_true) and os.path.isfile(video_test):
|
|
215
|
+
video_true_frames, video_test_frames, valid_frames = (
|
|
216
|
+
_fetch_video_frames(
|
|
217
|
+
video_true=video_true,
|
|
218
|
+
video_test=video_test,
|
|
219
|
+
)
|
|
220
|
+
)
|
|
221
|
+
elif os.path.isdir(video_true) and os.path.isdir(video_test):
|
|
222
|
+
# Glob videos
|
|
223
|
+
video_true_dir: pathlib.Path = pathlib.Path(video_true)
|
|
224
|
+
video_true_files = sorted(
|
|
225
|
+
[
|
|
226
|
+
file
|
|
227
|
+
for ext in _VIDEO_EXTENSIONS
|
|
228
|
+
for file in video_true_dir.rglob("*.{}".format(ext))
|
|
229
|
+
]
|
|
230
|
+
)
|
|
231
|
+
video_test_dir: pathlib.Path = pathlib.Path(video_test)
|
|
232
|
+
video_test_files = sorted(
|
|
233
|
+
[
|
|
234
|
+
file
|
|
235
|
+
for ext in _VIDEO_EXTENSIONS
|
|
236
|
+
for file in video_test_dir.rglob("*.{}".format(ext))
|
|
237
|
+
]
|
|
238
|
+
)
|
|
239
|
+
video_true_files = [file.as_posix() for file in video_true_files]
|
|
240
|
+
video_test_files = [file.as_posix() for file in video_test_files]
|
|
241
|
+
|
|
242
|
+
# select valid video files
|
|
243
|
+
video_true_files_selected = []
|
|
244
|
+
video_test_files_selected = []
|
|
245
|
+
for i in range(min(len(video_true_files), len(video_test_files))):
|
|
246
|
+
selected_video_true = video_true_files[i]
|
|
247
|
+
selected_video_test = video_test_files[i]
|
|
248
|
+
# Video pair must have the same basename
|
|
249
|
+
if os.path.basename(selected_video_test) == os.path.basename(
|
|
250
|
+
selected_video_true
|
|
251
|
+
):
|
|
252
|
+
video_true_files_selected.append(selected_video_true)
|
|
253
|
+
video_test_files_selected.append(selected_video_test)
|
|
254
|
+
|
|
255
|
+
video_true_files = video_true_files_selected.copy()
|
|
256
|
+
video_test_files = video_test_files_selected.copy()
|
|
257
|
+
if len(video_true_files) == 0:
|
|
258
|
+
logger.error(
|
|
259
|
+
"No valid Video pairs, please note that Video "
|
|
260
|
+
"pairs must have the same basename."
|
|
261
|
+
)
|
|
262
|
+
return None, None
|
|
263
|
+
logger.debug(f"video_true_files: {video_true_files}")
|
|
264
|
+
logger.debug(f"video_test_files: {video_test_files}")
|
|
265
|
+
|
|
266
|
+
# Fetch all frames
|
|
267
|
+
video_true_frames = []
|
|
268
|
+
video_test_frames = []
|
|
269
|
+
valid_frames = 0
|
|
270
|
+
|
|
271
|
+
for video_true_, video_test_ in zip(video_true_files, video_test_files):
|
|
272
|
+
video_true_frames_, video_test_frames_, valid_frames_ = (
|
|
273
|
+
_fetch_video_frames(
|
|
274
|
+
video_true=video_true_, video_test=video_test_
|
|
275
|
+
)
|
|
276
|
+
)
|
|
277
|
+
video_true_frames.extend(video_true_frames_)
|
|
278
|
+
video_test_frames.extend(video_test_frames_)
|
|
279
|
+
valid_frames += valid_frames_
|
|
280
|
+
else:
|
|
281
|
+
raise ValueError("video_true and video_test must be files or dirs.")
|
|
282
|
+
|
|
283
|
+
if valid_frames <= 0:
|
|
284
|
+
logger.debug("No valid frames to compare")
|
|
285
|
+
return None, None
|
|
286
|
+
|
|
287
|
+
total_metric = 0.0
|
|
288
|
+
valid_frames = 0 # reset
|
|
289
|
+
for frame1, frame2 in tqdm(
|
|
290
|
+
zip(video_true_frames, video_test_frames),
|
|
291
|
+
total=len(video_true_frames),
|
|
292
|
+
disable=DISABLE_VERBOSE,
|
|
293
|
+
):
|
|
294
|
+
metric = compute_frame_func(frame1, frame2)
|
|
295
|
+
if metric != float("inf"):
|
|
296
|
+
total_metric += metric
|
|
297
|
+
valid_frames += 1
|
|
298
|
+
|
|
299
|
+
if valid_frames > 0:
|
|
300
|
+
average_metric = total_metric / valid_frames
|
|
301
|
+
logger.debug(f"Average: {average_metric:.2f}")
|
|
302
|
+
return average_metric, valid_frames
|
|
303
|
+
else:
|
|
304
|
+
logger.debug("No valid frames to compare")
|
|
305
|
+
return None, None
|
|
306
|
+
|
|
307
|
+
|
|
308
|
+
compute_psnr = partial(
|
|
309
|
+
compute_dir_metric,
|
|
310
|
+
compute_file_func=compute_psnr_file,
|
|
311
|
+
)
|
|
312
|
+
|
|
313
|
+
compute_ssim = partial(
|
|
314
|
+
compute_dir_metric,
|
|
315
|
+
compute_file_func=compute_ssim_file,
|
|
316
|
+
)
|
|
317
|
+
|
|
318
|
+
compute_mse = partial(
|
|
319
|
+
compute_dir_metric,
|
|
320
|
+
compute_file_func=compute_mse_file,
|
|
321
|
+
)
|
|
322
|
+
|
|
323
|
+
compute_video_psnr = partial(
|
|
324
|
+
compute_video_metric,
|
|
325
|
+
compute_frame_func=compute_psnr_file,
|
|
326
|
+
)
|
|
327
|
+
compute_video_ssim = partial(
|
|
328
|
+
compute_video_metric,
|
|
329
|
+
compute_frame_func=compute_ssim_file,
|
|
330
|
+
)
|
|
331
|
+
compute_video_mse = partial(
|
|
332
|
+
compute_video_metric,
|
|
333
|
+
compute_frame_func=compute_mse_file,
|
|
334
|
+
)
|
|
335
|
+
|
|
336
|
+
|
|
337
|
+
# Entrypoints
|
|
338
|
+
def get_args():
|
|
339
|
+
parser = argparse.ArgumentParser(
|
|
340
|
+
description="CacheDiT's Metrics CLI",
|
|
341
|
+
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
342
|
+
)
|
|
343
|
+
METRICS_CHOICES = [
|
|
344
|
+
"psnr",
|
|
345
|
+
"ssim",
|
|
346
|
+
"mse",
|
|
347
|
+
"fid",
|
|
348
|
+
"all",
|
|
349
|
+
]
|
|
350
|
+
parser.add_argument(
|
|
351
|
+
"metric",
|
|
352
|
+
type=str,
|
|
353
|
+
default="psnr",
|
|
354
|
+
choices=METRICS_CHOICES,
|
|
355
|
+
help=f"Metric choices: {METRICS_CHOICES}",
|
|
356
|
+
)
|
|
357
|
+
parser.add_argument(
|
|
358
|
+
"--img-true",
|
|
359
|
+
"-i1",
|
|
360
|
+
type=str,
|
|
361
|
+
default=None,
|
|
362
|
+
help="Path to ground truth image or Dir to ground truth images",
|
|
363
|
+
)
|
|
364
|
+
parser.add_argument(
|
|
365
|
+
"--img-test",
|
|
366
|
+
"-i2",
|
|
367
|
+
type=str,
|
|
368
|
+
default=None,
|
|
369
|
+
help="Path to predicted image or Dir to predicted images",
|
|
370
|
+
)
|
|
371
|
+
parser.add_argument(
|
|
372
|
+
"--video-true",
|
|
373
|
+
"-v1",
|
|
374
|
+
type=str,
|
|
375
|
+
default=None,
|
|
376
|
+
help="Path to ground truth video or Dir to ground truth videos",
|
|
377
|
+
)
|
|
378
|
+
parser.add_argument(
|
|
379
|
+
"--video-test",
|
|
380
|
+
"-v2",
|
|
381
|
+
type=str,
|
|
382
|
+
default=None,
|
|
383
|
+
help="Path to predicted video or Dir to predicted videos",
|
|
384
|
+
)
|
|
385
|
+
parser.add_argument(
|
|
386
|
+
"--enable-verbose",
|
|
387
|
+
"-verbose",
|
|
388
|
+
action="store_true",
|
|
389
|
+
default=False,
|
|
390
|
+
help="Show metrics progress verbose",
|
|
391
|
+
)
|
|
392
|
+
return parser.parse_args()
|
|
393
|
+
|
|
394
|
+
|
|
395
|
+
def entrypoint():
|
|
396
|
+
args = get_args()
|
|
397
|
+
logger.debug(args)
|
|
398
|
+
|
|
399
|
+
if args.enable_verbose:
|
|
400
|
+
global DISABLE_VERBOSE
|
|
401
|
+
set_metrics_verbose(True)
|
|
402
|
+
DISABLE_VERBOSE = not get_metrics_verbose()
|
|
403
|
+
|
|
404
|
+
if args.img_true is not None and args.img_test is not None:
|
|
405
|
+
if any(
|
|
406
|
+
(
|
|
407
|
+
not os.path.exists(args.img_true),
|
|
408
|
+
not os.path.exists(args.img_test),
|
|
409
|
+
)
|
|
410
|
+
):
|
|
411
|
+
return
|
|
412
|
+
# img_true and img_test can be files or dirs
|
|
413
|
+
if args.metric == "psnr" or args.metric == "all":
|
|
414
|
+
img_psnr, n = compute_psnr(args.img_true, args.img_test)
|
|
415
|
+
logger.info(
|
|
416
|
+
f"{args.img_true} vs {args.img_test}, Num: {n}, PSNR: {img_psnr}"
|
|
417
|
+
)
|
|
418
|
+
if args.metric == "ssim" or args.metric == "all":
|
|
419
|
+
img_ssim, n = compute_ssim(args.img_true, args.img_test)
|
|
420
|
+
logger.info(
|
|
421
|
+
f"{args.img_true} vs {args.img_test}, Num: {n}, SSIM: {img_ssim}"
|
|
422
|
+
)
|
|
423
|
+
if args.metric == "mse" or args.metric == "all":
|
|
424
|
+
img_mse, n = compute_mse(args.img_true, args.img_test)
|
|
425
|
+
logger.info(
|
|
426
|
+
f"{args.img_true} vs {args.img_test}, Num: {n}, MSE: {img_mse}"
|
|
427
|
+
)
|
|
428
|
+
if args.metric == "fid" or args.metric == "all":
|
|
429
|
+
FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
|
|
430
|
+
img_fid, n = FID.compute_fid(args.img_true, args.img_test)
|
|
431
|
+
logger.info(
|
|
432
|
+
f"{args.img_true} vs {args.img_test}, Num: {n}, FID: {img_fid}"
|
|
433
|
+
)
|
|
434
|
+
if args.video_true is not None and args.video_test is not None:
|
|
435
|
+
if any(
|
|
436
|
+
(
|
|
437
|
+
not os.path.exists(args.video_true),
|
|
438
|
+
not os.path.exists(args.video_test),
|
|
439
|
+
)
|
|
440
|
+
):
|
|
441
|
+
return
|
|
442
|
+
# video_true and video_test can be files or dirs
|
|
443
|
+
if args.metric == "psnr" or args.metric == "all":
|
|
444
|
+
video_psnr, n = compute_video_psnr(args.video_true, args.video_test)
|
|
445
|
+
logger.info(
|
|
446
|
+
f"{args.video_true} vs {args.video_test}, Frames: {n}, PSNR: {video_psnr}"
|
|
447
|
+
)
|
|
448
|
+
if args.metric == "ssim" or args.metric == "all":
|
|
449
|
+
video_ssim, n = compute_video_ssim(args.video_true, args.video_test)
|
|
450
|
+
logger.info(
|
|
451
|
+
f"{args.video_true} vs {args.video_test}, Frames: {n}, SSIM: {video_ssim}"
|
|
452
|
+
)
|
|
453
|
+
if args.metric == "mse" or args.metric == "all":
|
|
454
|
+
video_mse, n = compute_video_mse(args.video_true, args.video_test)
|
|
455
|
+
logger.info(
|
|
456
|
+
f"{args.video_true} vs {args.video_test}, Frames: {n}, MSE: {video_mse}"
|
|
457
|
+
)
|
|
458
|
+
if args.metric == "fid" or args.metric == "all":
|
|
459
|
+
FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
|
|
460
|
+
video_fid, n = FID.compute_video_fid(
|
|
461
|
+
args.video_true, args.video_test
|
|
462
|
+
)
|
|
463
|
+
logger.info(
|
|
464
|
+
f"{args.video_true} vs {args.video_test}, Frames: {n}, FID: {video_fid}"
|
|
465
|
+
)
|
|
466
|
+
|
|
467
|
+
|
|
468
|
+
if __name__ == "__main__":
|
|
469
|
+
entrypoint()
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cache_dit
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.7
|
|
4
4
|
Summary: 🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration Toolbox for Diffusion Transformers
|
|
5
5
|
Author: DefTruth, vipshop.com, etc.
|
|
6
6
|
Maintainer: DefTruth, vipshop.com, etc
|
|
@@ -13,6 +13,8 @@ Requires-Dist: packaging
|
|
|
13
13
|
Requires-Dist: torch>=2.5.1
|
|
14
14
|
Requires-Dist: transformers>=4.51.3
|
|
15
15
|
Requires-Dist: diffusers>=0.33.1
|
|
16
|
+
Requires-Dist: scikit-image
|
|
17
|
+
Requires-Dist: scipy
|
|
16
18
|
Provides-Extra: all
|
|
17
19
|
Provides-Extra: dev
|
|
18
20
|
Requires-Dist: pre-commit; extra == "dev"
|
|
@@ -28,6 +30,8 @@ Requires-Dist: protobuf; extra == "dev"
|
|
|
28
30
|
Requires-Dist: sentencepiece; extra == "dev"
|
|
29
31
|
Requires-Dist: opencv-python-headless; extra == "dev"
|
|
30
32
|
Requires-Dist: ftfy; extra == "dev"
|
|
33
|
+
Requires-Dist: scikit-image; extra == "dev"
|
|
34
|
+
Requires-Dist: pytorch-fid; extra == "dev"
|
|
31
35
|
Dynamic: license-file
|
|
32
36
|
Dynamic: provides-extra
|
|
33
37
|
Dynamic: requires-dist
|
|
@@ -159,6 +163,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
|
|
|
159
163
|
- [⚡️Dynamic Block Prune](#dbprune)
|
|
160
164
|
- [🎉Context Parallelism](#context-parallelism)
|
|
161
165
|
- [🔥Torch Compile](#compile)
|
|
166
|
+
- [⚙️Metrics CLI](#metrics)
|
|
162
167
|
- [👋Contribute](#contribute)
|
|
163
168
|
- [©️License](#license)
|
|
164
169
|
|
|
@@ -500,6 +505,37 @@ torch._dynamo.config.accumulated_recompile_limit = 2048 # default is 256
|
|
|
500
505
|
|
|
501
506
|
Please check [bench.py](./bench/bench.py) for more details.
|
|
502
507
|
|
|
508
|
+
|
|
509
|
+
## ⚙️Metrics CLI
|
|
510
|
+
|
|
511
|
+
<div id="metrics"></div>
|
|
512
|
+
|
|
513
|
+
You can utilize the APIs provided by CacheDiT to quickly evaluate the accuracy losses caused by different cache configurations. For example:
|
|
514
|
+
|
|
515
|
+
```python
|
|
516
|
+
from cache_dit.metrics import compute_psnr
|
|
517
|
+
from cache_dit.metrics import compute_video_psnr
|
|
518
|
+
from cache_dit.metrics import FrechetInceptionDistance # FID
|
|
519
|
+
|
|
520
|
+
FID = FrechetInceptionDistance()
|
|
521
|
+
image_psnr, n = compute_psnr("true.png", "test.png") # Num: n
|
|
522
|
+
image_fid, n = FID.compute_fid("true_dir", "test_dir")
|
|
523
|
+
video_psnr, n = compute_video_psnr("true.mp4", "test.mp4") # Frames: n
|
|
524
|
+
```
|
|
525
|
+
|
|
526
|
+
Please check [test_metrics.py](./tests/test_metrics.py) for more details. Or, you can use `cache-dit-metrics-cli` tool. For examples:
|
|
527
|
+
|
|
528
|
+
```bash
|
|
529
|
+
cache-dit-metrics-cli -h # show usage
|
|
530
|
+
# all: PSNR, FID, SSIM, MSE, ..., etc.
|
|
531
|
+
cache-dit-metrics-cli all -i1 true.png -i2 test.png # image
|
|
532
|
+
cache-dit-metrics-cli all -i1 true_dir -i2 test_dir # image dir
|
|
533
|
+
cache-dit-metrics-cli all -v1 true.mp4 -v2 test.mp4 # video
|
|
534
|
+
cache-dit-metrics-cli all -v1 true_dir -v2 test_dir # video dir
|
|
535
|
+
cache-dit-metrics-cli fid -i1 true_dir -i2 test_dir # FID
|
|
536
|
+
cache-dit-metrics-cli psnr -i1 true_dir -i2 test_dir # PSNR
|
|
537
|
+
```
|
|
538
|
+
|
|
503
539
|
## 👋Contribute
|
|
504
540
|
<div id="contribute"></div>
|
|
505
541
|
|
|
@@ -1,17 +1,17 @@
|
|
|
1
1
|
cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
2
|
-
cache_dit/_version.py,sha256=
|
|
2
|
+
cache_dit/_version.py,sha256=Xk20v7uvkFqkpy9aLJzVngs1eKQn0FYUP2oyA1MEQUU,511
|
|
3
3
|
cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
|
|
4
4
|
cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
|
|
5
5
|
cache_dit/cache_factory/__init__.py,sha256=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
|
|
6
6
|
cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeDYnEqddA8,3987
|
|
7
7
|
cache_dit/cache_factory/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
8
8
|
cache_dit/cache_factory/dual_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
9
|
-
cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=
|
|
9
|
+
cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=7kMk6hvMDi-m2HP1qlj4p6qJhzZPjJol6IBneVGDs3E,71396
|
|
10
10
|
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=krNAICf-aS3JLmSG8vOB9tpLa04uYRcABsC8PMbVUKY,1870
|
|
11
|
-
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=
|
|
12
|
-
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=
|
|
13
|
-
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=
|
|
14
|
-
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=
|
|
11
|
+
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=3xUjvDzor9AkBkDUc0N7kZqM86MIdajuigesnicNzXE,2260
|
|
12
|
+
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=cIsov6Pf0dRyddqkzTA2CU-jSDotof8LQr-HIoY9T9M,2615
|
|
13
|
+
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=SO4q39PQuQ5QVHy5Z-ubiKdstzvQPedONN2J5oiGUh0,9955
|
|
14
|
+
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=8W9m-WeEVE2ytYi9udKEA8Wtb0EnvP3eT2A1Tu-d29k,2252
|
|
15
15
|
cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=EREHM5E1wxnL-uRXRAEege4HXraRp1oD_r1Zx4CsiKk,2596
|
|
16
16
|
cache_dit/cache_factory/dynamic_block_prune/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
17
17
|
cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=so1wGdb8W0ATwrjv7E5IEZLPcobybaY1HJa6hBYlOOQ,34698
|
|
@@ -33,8 +33,14 @@ cache_dit/compile/__init__.py,sha256=DfMdPleFFGADXLsr7zXui8BTz_y9futY6rNmNdh9y7k
|
|
|
33
33
|
cache_dit/compile/utils.py,sha256=KU60xc474Anbj7Y_FLRFmNxEjVYLLXkhbtCLXO7o_Tc,3699
|
|
34
34
|
cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
35
35
|
cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
36
|
-
cache_dit
|
|
37
|
-
cache_dit
|
|
38
|
-
cache_dit
|
|
39
|
-
cache_dit
|
|
40
|
-
cache_dit
|
|
36
|
+
cache_dit/metrics/__init__.py,sha256=RaUhl5dieF40RqnizGzR30qoJJ9dyMUEADwgwMaMQrE,575
|
|
37
|
+
cache_dit/metrics/config.py,sha256=ieOgD9ayz722RjVzk24bSIqS2D6o7TZjGk8KeXV-OLQ,551
|
|
38
|
+
cache_dit/metrics/fid.py,sha256=9Ivtazl6mW0Bon2VXa-Ia5Xj2ewxRD3V1Qkd69zYM3Y,17066
|
|
39
|
+
cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR0,12759
|
|
40
|
+
cache_dit/metrics/metrics.py,sha256=tzAtG_-fM1xPIBfRVFIBupvOWYzIO3xDq29Vy5rOBWc,14730
|
|
41
|
+
cache_dit-0.2.7.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
|
|
42
|
+
cache_dit-0.2.7.dist-info/METADATA,sha256=S0C1VGcXoWVjxmfX_755xttdoE9J0toSuHCZW9xaUBM,27608
|
|
43
|
+
cache_dit-0.2.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
44
|
+
cache_dit-0.2.7.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
|
|
45
|
+
cache_dit-0.2.7.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
|
|
46
|
+
cache_dit-0.2.7.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|