Spaces:
Sleeping
Sleeping
File size: 1,674 Bytes
32b2aaa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from typing import Callable, TypeVar, overload
import matplotlib.pyplot as plt
import numpy as np
def save_mels(path, *, targ_mel, pred_mel, cond_mel):
n = 3 if cond_mel is None else 4
plt.figure(figsize=(10, n * 4))
i = 1
plt.subplot(n, 1, i)
plt.imshow(pred_mel, origin="lower", interpolation="none")
plt.title(f"Pred mel {pred_mel.shape}")
i += 1
plt.subplot(n, 1, i)
plt.imshow(targ_mel, origin="lower", interpolation="none")
plt.title(f"GT mel {targ_mel.shape}")
i += 1
plt.subplot(n, 1, i)
pred_mel = pred_mel[:, : targ_mel.shape[1]]
targ_mel = targ_mel[:, : pred_mel.shape[1]]
plt.imshow(np.abs(pred_mel - targ_mel), origin="lower", interpolation="none")
plt.title(f"Diff mel {pred_mel.shape}, mse={np.mean((pred_mel - targ_mel)**2):.4f}")
i += 1
if cond_mel is not None:
plt.subplot(n, 1, i)
plt.imshow(cond_mel, origin="lower", interpolation="none")
plt.title(f"Cond mel {cond_mel.shape}")
i += 1
plt.savefig(path, dpi=480)
plt.close()
T = TypeVar("T")
@overload
def tree_map(fn: Callable, x: list[T]) -> list[T]:
...
@overload
def tree_map(fn: Callable, x: tuple[T]) -> tuple[T]:
...
@overload
def tree_map(fn: Callable, x: dict[str, T]) -> dict[str, T]:
...
@overload
def tree_map(fn: Callable, x: T) -> T:
...
def tree_map(fn: Callable, x):
if isinstance(x, list):
x = [tree_map(fn, xi) for xi in x]
elif isinstance(x, tuple):
x = (tree_map(fn, xi) for xi in x)
elif isinstance(x, dict):
x = {k: tree_map(fn, v) for k, v in x.items()}
else:
x = fn(x)
return x
|