Wan2GP/shared/utils/audio_video.py

import subprocess
import tempfile, os
import ffmpeg
import struct
from typing import Any
import torchvision.transforms.functional as TF
import torch.nn.functional as F
import cv2
import tempfile
import imageio
import binascii
import torchvision
import torch
from PIL import Image
import os.path as osp
import json
import numpy as np
import soundfile as sf
import zlib
import re
from .hdr import hdr10_x265_params, hdr10_zscale_filter, iter_hdr_gbrpf32_frames, iter_video_chunks

from .video_decode import probe_video_stream_metadata, resolve_media_binary
from .video_codecs import SUPPORTED_VIDEO_CONTAINERS, get_imageio_codec_params, get_video_encode_args, validate_video_output_settings
from .virtual_media import get_virtual_media_entry, parse_virtual_media_path, strip_virtual_media_suffix

def _ffmpeg_binary():
    return resolve_media_binary("ffmpeg") or "ffmpeg"


def _ffprobe_binary():
    return resolve_media_binary("ffprobe") or "ffprobe"


def rand_name(length=8, suffix=''):
    name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
    if suffix:
        if not suffix.startswith('.'):
            suffix = '.' + suffix
        name += suffix
    return name


def _prepare_audio_array(audio_data):
    if torch.is_tensor(audio_data):
        audio_data = audio_data.detach().cpu().float().numpy()
    else:
        audio_data = np.asarray(audio_data, dtype=np.float32)
    if audio_data.ndim == 2 and audio_data.shape[0] <= 8 and audio_data.shape[1] > audio_data.shape[0]:
        audio_data = audio_data.T
    return audio_data


def write_wav_file(path, audio_data, sample_rate):
    audio_array = _prepare_audio_array(audio_data)
    sf.write(path, audio_array, int(sample_rate))
    return path


def resample_audio_array(audio_data, source_sample_rate, target_sample_rate):
    audio_array = np.asarray(audio_data, dtype=np.float32)
    source_sample_rate = int(source_sample_rate or 0)
    target_sample_rate = int(target_sample_rate or 0)
    if audio_array.size == 0 or source_sample_rate <= 0 or target_sample_rate <= 0 or source_sample_rate == target_sample_rate:
        return audio_array.astype(np.float32, copy=False)
    import torchaudio.functional as taF
    wave = torch.from_numpy(audio_array.T.copy() if audio_array.ndim == 2 else audio_array[None].copy()).to(dtype=torch.float32)
    resampled = taF.resample(wave, source_sample_rate, target_sample_rate).cpu().numpy()
    return (resampled.T if audio_array.ndim == 2 else resampled[0]).astype(np.float32, copy=False)


def append_sliding_window_audio(existing_audio_data, existing_audio_path, generated_audio, audio_sampling_rate, committed_audio_samples, existing_audio_sample_rate=None):
    generated_audio = np.asarray(generated_audio, dtype=np.float32)
    if generated_audio.size == 0:
        return generated_audio
    prefix_sample_rate = int(existing_audio_sample_rate or audio_sampling_rate)
    if existing_audio_data is not None:
        prefix_audio = np.asarray(existing_audio_data, dtype=np.float32)
    elif existing_audio_path:
        prefix_audio, prefix_sample_rate = sf.read(os.fspath(existing_audio_path), dtype="float32", always_2d=generated_audio.ndim == 2)
    else:
        return generated_audio
    if prefix_sample_rate != int(audio_sampling_rate):
        prefix_audio = resample_audio_array(prefix_audio, prefix_sample_rate, audio_sampling_rate)
    prefix_audio = prefix_audio[:max(0, int(committed_audio_samples))]
    if prefix_audio.size == 0:
        return generated_audio
    if prefix_audio.ndim != generated_audio.ndim:
        prefix_audio = prefix_audio[:, None] if prefix_audio.ndim == 1 else prefix_audio
        generated_audio = generated_audio[:, None] if generated_audio.ndim == 1 else generated_audio
    if prefix_audio.ndim == 2 and prefix_audio.shape[1] != generated_audio.shape[1]:
        prefix_audio = np.repeat(prefix_audio[:, :1], generated_audio.shape[1], axis=1) if prefix_audio.shape[1] == 1 else prefix_audio[:, :generated_audio.shape[1]]
    return np.concatenate([prefix_audio, generated_audio], axis=0)


def create_silent_wav_file(output_dir=None, duration_seconds=0.0, sample_rate=16000, prefix="null_audio_"):
    sample_rate = int(sample_rate)
    num_samples = max(1, int(np.ceil(float(duration_seconds) * sample_rate)))
    fd, path = tempfile.mkstemp(prefix=prefix, suffix=".wav", dir=output_dir)
    os.close(fd)
    return write_wav_file(path, np.zeros(num_samples, dtype=np.float32), sample_rate)


def _compute_active_abs_amplitude(audio_data, active_mask=None):
    audio_data = np.asarray(audio_data, dtype=np.float32)
    if active_mask is not None:
        active_mask = np.asarray(active_mask, dtype=np.float32).reshape(-1) > 0.5
        if audio_data.ndim == 1:
            active_mask = active_mask[:audio_data.shape[0]]
            audio_data = audio_data[:active_mask.shape[0]][active_mask]
        else:
            active_mask = active_mask[:audio_data.shape[0]]
            audio_data = audio_data[:active_mask.shape[0]][active_mask]
    abs_audio = np.abs(audio_data).reshape(-1)
    if abs_audio.size == 0:
        return 0.0, 0.0
    avg_abs = float(abs_audio.mean())
    if avg_abs <= 0.0:
        return 0.0, 0.0
    threshold = 0.1 * avg_abs
    active_mask = abs_audio > threshold
    active_avg_abs = float(abs_audio[active_mask].mean()) if np.any(active_mask) else avg_abs
    return avg_abs, active_avg_abs


def normalize_audio_pair_volumes(audio1, audio2, active_mask1=None, active_mask2=None):
    audio1 = np.asarray(audio1, dtype=np.float32)
    audio2 = np.asarray(audio2, dtype=np.float32)
    avg1, active1 = _compute_active_abs_amplitude(audio1, active_mask1)
    avg2, active2 = _compute_active_abs_amplitude(audio2, active_mask2)
    midpoint = 0.5 * (active1 + active2)
    eps = 1e-8
    gain1 = midpoint / active1 if active1 > eps else 1.0
    gain2 = midpoint / active2 if active2 > eps else 1.0
    stats = {
        "audio1_avg_abs": float(avg1),
        "audio2_avg_abs": float(avg2),
        "audio1_active_avg_abs": float(active1),
        "audio2_active_avg_abs": float(active2),
        "target_active_avg_abs": float(midpoint),
        "audio1_gain": float(gain1),
        "audio2_gain": float(gain2),
    }
    return np.clip(audio1 * float(gain1), -1.0, 1.0), np.clip(audio2 * float(gain2), -1.0, 1.0), stats


def normalize_audio_pair_volumes_to_temp_files(audio_path1, audio_path2, output_dir=None, prefix="audio_norm_", active_mask1=None, active_mask2=None):
    audio1, sr1 = sf.read(os.fspath(audio_path1), dtype="float32", always_2d=False)
    audio2, sr2 = sf.read(os.fspath(audio_path2), dtype="float32", always_2d=False)
    norm1, norm2, stats = normalize_audio_pair_volumes(audio1, audio2, active_mask1=active_mask1, active_mask2=active_mask2)

    if output_dir is not None:
        os.makedirs(output_dir, exist_ok=True)

    fd1, out1 = tempfile.mkstemp(prefix=prefix + "1_", suffix=".wav", dir=output_dir)
    os.close(fd1)
    fd2, out2 = tempfile.mkstemp(prefix=prefix + "2_", suffix=".wav", dir=output_dir)
    os.close(fd2)
    sf.write(out1, norm1, int(sr1))
    sf.write(out2, norm2, int(sr2))
    return out1, out2, stats


def _get_audio_codec_settings(codec_key):
    if not codec_key:
        codec_key = "wav"
    codec_key = str(codec_key).lower()
    if codec_key == "mp3":
        codec_key = "mp3_192"
    settings = {
        "wav": {"ext": "wav", "format": "wav"},
        "mp3_128": {"ext": "mp3", "format": "mp3", "bitrate": "128k"},
        "mp3_192": {"ext": "mp3", "format": "mp3", "bitrate": "192k"},
        "mp3_320": {"ext": "mp3", "format": "mp3", "bitrate": "320k"},
    }
    return settings.get(codec_key, settings["wav"])


def get_mp4_audio_codec_settings(codec_key):
    codec_key = "aac_128" if not codec_key else str(codec_key).lower()
    settings = {
        "aac_128": {"codec": "aac", "bitrate": "128k", "ext": ".aac"},
        "aac_192": {"codec": "aac", "bitrate": "192k", "ext": ".aac"},
        "aac_256": {"codec": "aac", "bitrate": "256k", "ext": ".aac"},
        "aac_320": {"codec": "aac", "bitrate": "320k", "ext": ".aac"},
        "alac": {"codec": "alac", "bitrate": None, "ext": ".m4a"},
    }
    return settings.get(codec_key, settings["aac_128"])


def _infer_video_dimensions(tensor):
    if torch.is_tensor(tensor):
        if tensor.ndim == 5:
            return int(tensor.shape[-1]), int(tensor.shape[-2])
        if tensor.ndim == 4:
            if tensor.shape[-1] in (1, 3, 4):
                return int(tensor.shape[2]), int(tensor.shape[1])
            return int(tensor.shape[-1]), int(tensor.shape[-2])
    if isinstance(tensor, (list, tuple)):
        for chunk in tensor:
            dims = _infer_video_dimensions(chunk)
            if dims is not None:
                return dims
    return None


def _validate_video_save_settings(codec_type, container, tensor):
    dims = _infer_video_dimensions(tensor)
    width = height = None
    if dims is not None:
        width, height = dims
    error = validate_video_output_settings(codec_type, container, width=width, height=height, allowed_containers=SUPPORTED_VIDEO_CONTAINERS)
    if error is not None:
        raise RuntimeError(error)


def _crf_from_video_codec(codec_key: str | None, default: str = "18") -> str:
    codec_key = str(codec_key or "").strip().lower()
    if re.fullmatch(r"\d+", codec_key):
        return codec_key
    match = re.search(r"_(\d+)$", codec_key)
    return match.group(1) if match is not None else str(default)


def get_hdr_video_encode_args(codec_key: str | None, container: str | None) -> list[str]:
    crf = _crf_from_video_codec(codec_key, default="18")
    return [
        "-vf", hdr10_zscale_filter(),
        "-c:v", "libx265",
        "-preset", "medium",
        "-crf", crf,
        "-pix_fmt", "yuv420p10le",
        "-tag:v", "hvc1",
        "-color_primaries", "bt2020",
        "-color_trc", "smpte2084",
        "-colorspace", "bt2020nc",
        "-x265-params", hdr10_x265_params(),
    ]


def get_audio_codec_extension(codec_key):
    return _get_audio_codec_settings(codec_key)["ext"]


def _run_ffmpeg_encode(input_path, output_path, codec, bitrate=None, sample_rate=None, drop_video=False):
    cmd = [_ffmpeg_binary(), "-y", "-v", "error", "-i", input_path]
    if drop_video:
        cmd.append("-vn")
    cmd += ["-c:a", codec]
    if bitrate:
        cmd += ["-b:a", bitrate]
    if sample_rate:
        cmd += ["-ar", str(int(sample_rate))]
    cmd.append(output_path)
    subprocess.run(cmd, check=True, capture_output=True, text=True)


def save_audio_file(path, audio_data, sample_rate, codec_key="wav"):
    settings = _get_audio_codec_settings(codec_key)
    ext = settings["ext"]
    if not path.lower().endswith(f".{ext}"):
        path = osp.splitext(path)[0] + f".{ext}"
    if settings["format"] == "wav":
        return write_wav_file(path, audio_data, sample_rate)
    fd, tmp_path = tempfile.mkstemp(suffix=".wav", prefix="audio_")
    os.close(fd)
    try:
        write_wav_file(tmp_path, audio_data, sample_rate)
        _run_ffmpeg_encode(tmp_path, path, "libmp3lame", bitrate=settings.get("bitrate"), sample_rate=sample_rate)
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass
    return path


def _resolve_virtual_audio_segment(video_path: str) -> tuple[str, dict[str, Any], int]:
    if isinstance(video_path, Image.Image):
        return "", {}, 0
    if get_virtual_media_entry(video_path) is not None:
        return "", {}, 0
    spec = parse_virtual_media_path(video_path)
    source_path = os.fspath(strip_virtual_media_suffix(video_path))
    time_args: dict[str, Any] = {}
    if spec is None:
        return source_path, time_args, 0
    metadata = probe_video_stream_metadata(video_path)
    if metadata is not None and metadata.get("virtual_end_frame") is not None:
        start_frame = int(metadata.get("virtual_start_frame") or 0)
        end_frame = int(metadata.get("virtual_end_frame") or start_frame)
        fps_float = float(metadata.get("fps_float") or metadata.get("fps") or 0.0)
        if fps_float > 0:
            time_args["ss"] = max(0.0, start_frame / fps_float)
            time_args["to"] = max(time_args["ss"], (end_frame + 1) / fps_float)
    audio_track_no = 1 if spec.audio_track_no is None else max(1, int(spec.audio_track_no))
    return source_path, time_args, audio_track_no - 1


def extract_audio_track_to_wav(video_path, output_path):
    if not video_path:
        return None
    if isinstance(video_path, Image.Image):
        return None
    video_path = os.fspath(video_path)
    source_path, time_args, audio_track_index = _resolve_virtual_audio_segment(video_path)
    if len(source_path) == 0:
        return None
    import ffmpeg
    try:
        output_kwargs = {"map": f"0:a:{audio_track_index}", "acodec": "pcm_s16le"}
        ffmpeg.input(source_path, **time_args).output(output_path, **output_kwargs).overwrite_output().run(cmd=_ffmpeg_binary(), quiet=True)
    except ffmpeg.Error as err:
        stderr = getattr(err, "stderr", b"")
        if isinstance(stderr, (bytes, bytearray)):
            stderr = stderr.decode("utf-8", errors="ignore")
        stderr = (stderr or str(err)).strip()
        raise RuntimeError(f"ffmpeg audio extract failed for {source_path} -> {output_path}: {stderr}") from err
    return output_path



def extract_audio_tracks(source_video, verbose=False, query_only=False, codec_key="aac_128", temp_format=None):
    """
    Extract all audio tracks from a source video into temporary audio files.

    Returns:
        Tuple:
          - List of temp file paths for extracted audio tracks
          - List of corresponding metadata dicts:
              {'codec', 'sample_rate', 'channels', 'duration', 'language'}
              where 'duration' is set to container duration (for consistency).
    """
    if isinstance(source_video, Image.Image):
        return 0 if query_only else ([], [])
    source_path, time_args, selected_track_index = _resolve_virtual_audio_segment(source_video)
    if len(source_path) == 0:
        return 0 if query_only else ([], [])
    if not os.path.exists(source_path):
        msg = f"ffprobe skipped; file not found: {source_video}"
        if verbose:
            print(msg)
        raise FileNotFoundError(msg)

    try:
        probe = ffmpeg.probe(source_path, cmd=_ffprobe_binary())
    except ffmpeg.Error as err:
        stderr = getattr(err, 'stderr', b'')
        if isinstance(stderr, (bytes, bytearray)):
            stderr = stderr.decode('utf-8', errors='ignore')
        stderr = (stderr or str(err)).strip()
        message = f"ffprobe failed for {source_path}: {stderr}"
        if verbose:
            print(message)
        raise RuntimeError(message) from err
    audio_streams = [s for s in probe['streams'] if s['codec_type'] == 'audio']
    container_duration = float(probe['format'].get('duration', 0.0))
    if selected_track_index is not None:
        audio_streams = [audio_streams[selected_track_index]] if 0 <= selected_track_index < len(audio_streams) else []

    if not audio_streams:
        if query_only: return 0
        if verbose: print(f"No audio track found in {source_video}")
        return [], []

    if query_only:
        return len(audio_streams)

    if verbose:
        print(f"Found {len(audio_streams)} audio track(s), container duration = {container_duration:.3f}s")

    file_paths = []
    metadata = []
    if temp_format == "wav":
        audio_settings = {"codec": "pcm_s16le", "bitrate": None, "ext": ".wav"}
    else:
        audio_settings = get_mp4_audio_codec_settings(codec_key)

    for i, stream in enumerate(audio_streams):
        fd, temp_path = tempfile.mkstemp(suffix=f'_track{i}{audio_settings["ext"]}', prefix='audio_')
        os.close(fd)

        file_paths.append(temp_path)
        metadata.append({
            'codec': stream.get('codec_name'),
            'sample_rate': int(stream.get('sample_rate', 0)),
            'channels': int(stream.get('channels', 0)),
            'duration': container_duration,
            'language': stream.get('tags', {}).get('language', None)
        })

        stream_index = i if selected_track_index is None else selected_track_index
        output_kwargs = {f'map': f'0:a:{stream_index}', 'acodec': audio_settings["codec"]}
        if audio_settings["bitrate"]:
            output_kwargs['b:a'] = audio_settings["bitrate"]
        ffmpeg.input(source_path, **time_args).output(temp_path, **output_kwargs).overwrite_output().run(cmd=_ffmpeg_binary(), quiet=not verbose)

    return file_paths, metadata



def combine_and_concatenate_video_with_audio_tracks(
    save_path_tmp, video_path,
    source_audio_tracks, new_audio_tracks,
    source_audio_duration, audio_sampling_rate,
    new_audio_from_start=False,
    source_audio_metadata=None,
    audio_codec_key="aac_128",
    verbose = False
):
    audio_settings = get_mp4_audio_codec_settings(audio_codec_key)
    audio_codec = audio_settings["codec"]
    audio_bitrate = audio_settings["bitrate"]
    inputs, filters, maps, idx = ['-i', video_path], [], ['-map', '0:v'], 1
    metadata_args = []
    sources = source_audio_tracks or []
    news = new_audio_tracks or []

    duplicate_source = len(sources) == 1 and len(news) > 1
    N = len(news) if source_audio_duration == 0 else max(len(sources), len(news)) or 1

    for i in range(N):
        s = (sources[i] if i < len(sources)
             else sources[0] if duplicate_source else None)
        n = news[i] if len(news) == N else (news[0] if news else None)

        if source_audio_duration == 0:
            if n:
                inputs += ['-i', n]
                filters.append(f'[{idx}:a]apad=pad_dur=100[aout{i}]')
                idx += 1
            else:
                filters.append(f'anullsrc=r={audio_sampling_rate}:cl=mono,apad=pad_dur=100[aout{i}]')
        else:
            if s:
                inputs += ['-i', s]
                meta = source_audio_metadata[i] if source_audio_metadata and i < len(source_audio_metadata) else {}
                needs_filter = (
                    meta.get('codec') != audio_codec or
                    meta.get('sample_rate') != audio_sampling_rate or
                    meta.get('channels') != 1 or
                    meta.get('duration', 0) < source_audio_duration
                )
                if needs_filter:
                    filters.append(
                        f'[{idx}:a]aresample={audio_sampling_rate},aformat=channel_layouts=mono,'
                        f'apad=pad_dur={source_audio_duration},atrim=0:{source_audio_duration},asetpts=PTS-STARTPTS[s{i}]')
                else:
                    filters.append(
                        f'[{idx}:a]apad=pad_dur={source_audio_duration},atrim=0:{source_audio_duration},asetpts=PTS-STARTPTS[s{i}]')
                if lang := meta.get('language'):
                    metadata_args += ['-metadata:s:a:' + str(i), f'language={lang}']
                idx += 1
            else:
                filters.append(
                    f'anullsrc=r={audio_sampling_rate}:cl=mono,atrim=0:{source_audio_duration},asetpts=PTS-STARTPTS[s{i}]')

            if n:
                inputs += ['-i', n]
                start = '0' if new_audio_from_start else source_audio_duration
                filters.append(
                    f'[{idx}:a]aresample={audio_sampling_rate},aformat=channel_layouts=mono,'
                    f'atrim=start={start},asetpts=PTS-STARTPTS[n{i}]')
                filters.append(f'[s{i}][n{i}]concat=n=2:v=0:a=1[aout{i}]')
                idx += 1
            else:
                filters.append(f'[s{i}]apad=pad_dur=100[aout{i}]')

        maps += ['-map', f'[aout{i}]']

    cmd = [_ffmpeg_binary(), '-y', *inputs,
           '-filter_complex', ';'.join(filters),  # ✅ Only change made
           *maps, *metadata_args,
           '-c:v', 'copy',
           '-c:a', audio_codec,
           '-ar', str(audio_sampling_rate),
           '-ac', '1',
           '-shortest', save_path_tmp]
    if audio_bitrate:
        cmd[-6:-6] = ['-b:a', audio_bitrate]

    if verbose:
        print(f"ffmpeg command: {cmd}")
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        raise Exception(f"FFmpeg error: {e.stderr}")


def combine_video_with_audio_tracks(target_video, audio_tracks, output_video,
                                     audio_metadata=None, audio_codec_key="aac_128", verbose=False):
    if not audio_tracks:
        if verbose: print("No audio tracks to combine."); return False

    dur = float(next(s for s in ffmpeg.probe(target_video, cmd=_ffprobe_binary())['streams']
                     if s['codec_type'] == 'video')['duration'])
    if verbose: print(f"Video duration: {dur:.3f}s")

    cmd = [_ffmpeg_binary(), '-y', '-i', target_video]
    for path in audio_tracks:
        cmd += ['-i', path]

    cmd += ['-map', '0:v']
    for i in range(len(audio_tracks)):
        cmd += ['-map', f'{i+1}:a']

    for i, meta in enumerate(audio_metadata or []):
        if (lang := meta.get('language')):
            cmd += ['-metadata:s:a:' + str(i), f'language={lang}']

    audio_settings = get_mp4_audio_codec_settings(audio_codec_key)
    cmd += ['-c:v', 'copy', '-c:a', audio_settings["codec"]]
    if audio_settings["bitrate"]:
        cmd += ['-b:a', audio_settings["bitrate"]]
    cmd += ['-t', str(dur), output_video]

    result = subprocess.run(cmd, capture_output=not verbose, text=True)
    if result.returncode != 0:
        raise Exception(f"FFmpeg error:\n{result.stderr}")
    if verbose:
        print(f"Created {output_video} with {len(audio_tracks)} audio track(s)")
    return True


def cleanup_temp_audio_files(audio_tracks, verbose=False):
    """
    Clean up temporary audio files.
    
    Args:
        audio_tracks: List of audio file paths to delete
        verbose: Enable verbose output (default: False)
        
    Returns:
        Number of files successfully deleted
    """
    deleted_count = 0
    
    for audio_path in audio_tracks:
        try:
            if os.path.exists(audio_path):
                os.unlink(audio_path)
                deleted_count += 1
                if verbose:
                    print(f"Cleaned up {audio_path}")
        except PermissionError:
            print(f"Warning: Could not delete {audio_path} (file may be in use)")
        except Exception as e:
            print(f"Warning: Error deleting {audio_path}: {e}")
    
    if verbose and deleted_count > 0:
        print(f"Successfully deleted {deleted_count} temporary audio file(s)")
    
    return deleted_count


def save_video(tensor,
                save_file=None,
                fps=30,
                codec_type='libx264_8',
                container='mp4',
                nrow=8,
                normalize=True,
                value_range=(-1, 1),
                retry=5):
    """Save tensor as video with configurable codec and container options."""
        
    if torch.is_tensor(tensor) and len(tensor.shape) == 4:
        tensor = tensor.unsqueeze(0)

    _validate_video_save_settings(codec_type, container, tensor)
        
    suffix = f'.{container}'
    cache_file = osp.join('/tmp', rand_name(suffix=suffix)) if save_file is None else save_file
    if not cache_file.endswith(suffix):
        cache_file = osp.splitext(cache_file)[0] + suffix
    
    # Configure codec parameters
    codec_params = _get_codec_params(codec_type, container)
    
    # Process and save
    error = None
    for _ in range(retry):
        try:
            # Write video (silence ffmpeg logs)
            writer = imageio.get_writer(cache_file, fps=fps, ffmpeg_log_level='error', **codec_params)
            try:
                if torch.is_tensor(tensor):
                    # Stream frames to avoid materializing the full video on CPU.
                    if tensor.dtype == torch.uint8 and tensor.ndim == 5 and tensor.shape[0] == 1 and nrow == 1:
                        frames = tensor[0].permute(1, 2, 3, 0)
                        for frame in frames:
                            writer.append_data(frame.cpu().numpy())
                    else:
                        if tensor.dtype == torch.uint8:
                            tensor = tensor.float().div_(127.5).sub_(1.0)
                        for u in tensor.unbind(2):
                            u = u.clamp(min(value_range), max(value_range))
                            grid = torchvision.utils.make_grid(
                                u, nrow=nrow, normalize=normalize, value_range=value_range
                            )
                            frame = grid.mul(255).type(torch.uint8).permute(1, 2, 0).cpu().numpy()
                            writer.append_data(frame)
                elif isinstance(tensor, (list, tuple)) and tensor and torch.is_tensor(tensor[0]):
                    for chunk in tensor:
                        if chunk is None:
                            continue
                        if chunk.ndim == 4:
                            if chunk.shape[-1] in (1, 3, 4):
                                frames = chunk
                            else:
                                frames = chunk.permute(1, 2, 3, 0)
                            for frame in frames:
                                writer.append_data(frame.cpu().numpy())
                        else:
                            writer.append_data(chunk)
                else:
                    for frame in tensor:
                        writer.append_data(frame)
            finally:
                writer.close()

            return cache_file

        except Exception as e:
            error = e
            print(f"error saving {save_file}: {e}")


def save_hdr_video(
                tensor,
                save_file=None,
                fps=30,
                codec_type='libx264_8',
                container='mp4',
                preview_exposure=0.0,
                retry=5):
    """Save linear HDR video as a tagged 10-bit HEVC HDR file."""
    suffix = f'.{container}'
    output_file = osp.join('/tmp', rand_name(suffix=suffix)) if save_file is None else save_file
    if not output_file.endswith(suffix):
        output_file = osp.splitext(output_file)[0] + suffix
    ffmpeg_path = resolve_media_binary("ffmpeg")
    if ffmpeg_path is None:
        raise RuntimeError("ffmpeg binary not found")

    width = height = None
    for chunk in iter_video_chunks(tensor):
        if chunk is None:
            continue
        cur = chunk[0] if chunk.ndim == 5 and chunk.shape[0] == 1 else chunk
        if cur.ndim == 4:
            height, width = int(cur.shape[2]), int(cur.shape[3])
            break
    if width is None or height is None:
        raise RuntimeError("Unable to determine HDR video dimensions.")

    error = None
    for _ in range(retry):
        cmd = [
            ffmpeg_path, "-y", "-v", "error",
            "-f", "rawvideo",
            "-pix_fmt", "gbrpf32le",
            "-video_size", f"{width}x{height}",
            "-framerate", f"{float(fps):.12g}",
            "-i", "pipe:0",
            *get_hdr_video_encode_args(codec_type, container),
            "-an",
            output_file,
        ]
        process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        try:
            assert process.stdin is not None
            wrote_frame = False
            for frame_bytes in iter_hdr_gbrpf32_frames(tensor):
                process.stdin.write(frame_bytes)
                wrote_frame = True
            if not wrote_frame:
                raise RuntimeError("No HDR frames available to save.")
            process.stdin.close()
            stderr = process.stderr.read().decode("utf-8", errors="ignore").strip() if process.stderr is not None else ""
            ret = process.wait()
            if ret != 0:
                raise RuntimeError(stderr or "ffmpeg HDR encode failed")
            return output_file
        except Exception as e:
            error = e
            try:
                if process.stdin is not None and not process.stdin.closed:
                    process.stdin.close()
            except Exception:
                pass
            process.kill()
            print(f"error saving HDR {save_file}: {e}")
    raise error or RuntimeError(f"Failed to save HDR video: {save_file}")


def _get_codec_params(codec_type, container):
    """Get codec parameters based on codec type and container."""
    return get_imageio_codec_params(codec_type, container)




def save_image(tensor,
                save_file,
                nrow=8,
                normalize=True,
                value_range=(-1, 1),
                quality='jpeg_95',  # 'jpeg_95', 'jpeg_85', 'jpeg_70', 'jpeg_50', 'webp_95', 'webp_85', 'webp_70', 'webp_50', 'png', 'webp_lossless'
                retry=5):
    """Save tensor as image with configurable format and quality."""

    RGBA = tensor.shape[0] == 4
    if RGBA:
        quality = "png"

    # Get format and quality settings
    format_info = _get_format_info(quality)
    
    # Rename file extension to match requested format
    save_file = osp.splitext(save_file)[0] + format_info['ext']
    
    # Save image
    error = None
                         
    for _ in range(retry):
        try:
            if format_info['use_pil'] or RGBA:
                # Use PIL for WebP and advanced options
                if tensor.dtype == torch.uint8:
                    grid = torchvision.utils.make_grid(tensor, nrow=nrow, normalize=False).permute(1, 2, 0).cpu().numpy()
                else:
                    tensor = tensor.clamp(min(value_range), max(value_range))
                    grid = torchvision.utils.make_grid(tensor, nrow=nrow, normalize=normalize, value_range=value_range)
                    grid = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
                mode = 'RGBA' if RGBA else 'RGB'
                img = Image.fromarray(grid, mode=mode)
                img.save(save_file, **format_info['params'])
            else:
                # Use torchvision for JPEG and PNG
                was_uint8 = tensor.dtype == torch.uint8
                tensor = tensor.float().div_(255.0) if was_uint8 else tensor.clamp(min(value_range), max(value_range))
                torchvision.utils.save_image(tensor, save_file, nrow=nrow, normalize=False if was_uint8 else normalize, value_range=value_range, **format_info['params'])
            break
        except Exception as e:
            error = e
            continue
    else:
        print(f'cache_image failed, error: {error}', flush=True)
    
    return save_file


def _get_format_info(quality):
    """Get format extension and parameters."""
    formats = {
        # JPEG with PIL (so 'quality' works)
        'jpeg_95': {'ext': '.jpg', 'params': {'quality': 95}, 'use_pil': True},
        'jpeg_85': {'ext': '.jpg', 'params': {'quality': 85}, 'use_pil': True},
        'jpeg_70': {'ext': '.jpg', 'params': {'quality': 70}, 'use_pil': True},
        'jpeg_50': {'ext': '.jpg', 'params': {'quality': 50}, 'use_pil': True},

        # PNG with torchvision
        'png': {'ext': '.png', 'params': {}, 'use_pil': False},

        # WebP with PIL (for quality control)
        'webp_95': {'ext': '.webp', 'params': {'quality': 95}, 'use_pil': True},
        'webp_85': {'ext': '.webp', 'params': {'quality': 85}, 'use_pil': True},
        'webp_70': {'ext': '.webp', 'params': {'quality': 70}, 'use_pil': True},
        'webp_50': {'ext': '.webp', 'params': {'quality': 50}, 'use_pil': True},
        'webp_lossless': {'ext': '.webp', 'params': {'lossless': True}, 'use_pil': True},
    }
    return formats.get(quality, formats['jpeg_95'])


from PIL import Image, PngImagePlugin

def _enc_uc(s):
    try: return b"ASCII\0\0\0" + s.encode("ascii")
    except UnicodeEncodeError: return b"UNICODE\0" + s.encode("utf-16le")

def _dec_uc(b):
    if not isinstance(b, (bytes, bytearray)):
        try: b = bytes(b)
        except Exception: return None
    if b.startswith(b"ASCII\0\0\0"): return b[8:].decode("ascii", "ignore")
    if b.startswith(b"UNICODE\0"):   return b[8:].decode("utf-16le", "ignore")
    return b.decode("utf-8", "ignore")


def _blank_exif_dict():
    return {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}


def _load_exif_dict(image_path, ext):
    import piexif
    try:
        if ext in (".jpg", ".jpeg"):
            return piexif.load(image_path)
        if ext == ".webp":
            with Image.open(image_path) as im:
                exif_bytes = im.info.get("exif")
            return piexif.load(exif_bytes) if exif_bytes else _blank_exif_dict()
    except Exception:
        pass
    return _blank_exif_dict()


def _insert_exif_user_comment(image_path, comment_text, ext):
    import piexif
    exif_dict = _load_exif_dict(image_path, ext)
    exif_dict.setdefault("Exif", {})
    exif_dict["Exif"][piexif.ExifIFD.UserComment] = _enc_uc(comment_text)
    piexif.insert(piexif.dump(exif_dict), image_path)


_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


def _build_png_chunk(chunk_type, data):
    return struct.pack(">I", len(data)) + chunk_type + data + struct.pack(">I", zlib.crc32(chunk_type + data) & 0xffffffff)


def _is_png_comment_chunk(chunk_type, data):
    if chunk_type not in {b"tEXt", b"zTXt", b"iTXt"}:
        return False
    return data.split(b"\x00", 1)[0] == b"comment"


def _write_png_comment_metadata(image_path, comment_text):
    raw = open(image_path, "rb").read()
    if not raw.startswith(_PNG_SIGNATURE):
        raise ValueError("Invalid PNG signature")
    comment_chunk = _build_png_chunk(b"iTXt", b"comment\x00\x00\x00\x00\x00" + comment_text.encode("utf-8"))
    out = bytearray(_PNG_SIGNATURE)
    pos = len(_PNG_SIGNATURE)
    inserted = False
    while pos < len(raw):
        if pos + 8 > len(raw):
            raise ValueError("Corrupted PNG chunk header")
        length = struct.unpack(">I", raw[pos:pos + 4])[0]
        chunk_type = raw[pos + 4:pos + 8]
        end = pos + 12 + length
        if end > len(raw):
            raise ValueError("Corrupted PNG chunk payload")
        chunk_data = raw[pos + 8:pos + 8 + length]
        chunk = raw[pos:end]
        pos = end
        if _is_png_comment_chunk(chunk_type, chunk_data):
            continue
        if not inserted and chunk_type == b"IDAT":
            out.extend(comment_chunk)
            inserted = True
        out.extend(chunk)
    if not inserted:
        raise ValueError("PNG image data chunk not found")
    with open(image_path, "wb") as writer:
        writer.write(out)

def save_image_metadata(image_path, metadata_dict, **save_kwargs):
    try:
        j = json.dumps(metadata_dict, ensure_ascii=False)
        ext = os.path.splitext(image_path)[1].lower()
        if ext == ".png":
            _write_png_comment_metadata(image_path, j); return True
        if ext in (".jpg", ".jpeg", ".webp"):
            _insert_exif_user_comment(image_path, j, ext); return True
        raise ValueError("Unsupported format")
    except Exception as e:
        print(f"Error saving metadata: {e}"); return False

def read_image_metadata(image_path):
    try:
        ext = os.path.splitext(image_path)[1].lower()
        with Image.open(image_path) as im:
            if ext == ".png":
                val = (getattr(im, "text", {}) or {}).get("comment") or im.info.get("comment")
                return json.loads(val) if val else None
            if ext in (".jpg", ".jpeg"):
                import piexif
                try:
                    uc = piexif.load(image_path).get("Exif", {}).get(piexif.ExifIFD.UserComment)
                    s = _dec_uc(uc) if uc else None
                    if s:
                        return json.loads(s)
                except Exception:
                    pass
                val = im.info.get("comment")
                if isinstance(val, (bytes, bytearray)): val = val.decode("utf-8", "ignore")
                if val:
                    try: return json.loads(val)
                    except Exception: pass
                exif = getattr(im, "getexif", lambda: None)()
                if exif:
                    uc = exif.get(37510)  # UserComment
                    s = _dec_uc(uc) if uc else None
                    if s:
                        try: return json.loads(s)
                        except Exception: pass
                return None
            if ext == ".webp":
                import piexif
                exif_bytes = im.info.get("exif")
                if not exif_bytes: return None
                uc = piexif.load(exif_bytes).get("Exif", {}).get(piexif.ExifIFD.UserComment)
                s = _dec_uc(uc) if uc else None
                return json.loads(s) if s else None
            return None
    except Exception as e:
        print(f"Error reading metadata: {e}"); return None