73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
"""Shared utilities for component tests.
|
|
|
|
Component tests run inside the Docker image against real GPU models. They
|
|
write their output artefacts (MP4s, PNGs, logs) to ``_out/`` so you can
|
|
visually inspect results.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
|
|
OUT_DIR = os.path.join(os.path.dirname(__file__), "_out")
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
|
|
# A tiny 256x256 portrait PNG lives next to the component tests so tests
|
|
# don't need a user-supplied file. If it's missing we synthesise one on
|
|
# the fly.
|
|
SAMPLE_AVATAR = os.path.join(os.path.dirname(__file__), "sample_avatar.png")
|
|
|
|
|
|
def get_logger(name: str) -> logging.Logger:
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(name)s %(levelname)s %(message)s",
|
|
stream=sys.stdout,
|
|
)
|
|
return logging.getLogger(name)
|
|
|
|
|
|
def ensure_sample_avatar() -> str:
|
|
"""Guarantee a usable avatar image exists. Returns its path."""
|
|
if os.path.isfile(SAMPLE_AVATAR):
|
|
return SAMPLE_AVATAR
|
|
# Synthesise a simple gradient PNG as a last resort (won't look like a
|
|
# person but is valid input for Wan2.2 so the pipeline doesn't fail).
|
|
try:
|
|
from PIL import Image # type: ignore[import-not-found]
|
|
except ImportError:
|
|
import imageio.v3 as iio # type: ignore[import-not-found]
|
|
arr = np.zeros((256, 256, 3), dtype=np.uint8)
|
|
for y in range(256):
|
|
arr[y, :, 0] = y
|
|
arr[y, :, 1] = 255 - y
|
|
arr[y, :, 2] = 128
|
|
iio.imwrite(SAMPLE_AVATAR, arr)
|
|
return SAMPLE_AVATAR
|
|
|
|
arr = np.zeros((256, 256, 3), dtype=np.uint8)
|
|
for y in range(256):
|
|
arr[y, :, 0] = y
|
|
arr[y, :, 1] = 255 - y
|
|
arr[y, :, 2] = 128
|
|
Image.fromarray(arr).save(SAMPLE_AVATAR)
|
|
return SAMPLE_AVATAR
|
|
|
|
|
|
def write_bytes(name: str, data: bytes) -> str:
|
|
"""Write an artefact to _out/<name> and return the full path."""
|
|
path = os.path.join(OUT_DIR, name)
|
|
with open(path, "wb") as f:
|
|
f.write(data)
|
|
return path
|
|
|
|
|
|
def synth_tone(seconds: float, sample_rate: int = 24000, freq: float = 220.0) -> np.ndarray:
|
|
"""Return a float32 sine tone usable as stand-in TTS audio."""
|
|
t = np.arange(int(seconds * sample_rate), dtype=np.float32) / sample_rate
|
|
return (0.2 * np.sin(2 * np.pi * freq * t)).astype(np.float32)
|