core_anim/space_config.py

from __future__ import annotations

import configparser
from dataclasses import dataclass
import os
from pathlib import Path


ROOT = Path(__file__).resolve().parent.parent


def _bool(value: str) -> bool:
    return value.strip().lower() in {"1", "true", "yes", "on"}


@dataclass(frozen=True, slots=True)
class SpaceConfig:
    width: int
    height: int
    seconds: float
    fps: int
    steps: int
    split_step: int
    cfg: float
    shift: float
    sampler: str
    scheduler: str
    start_strength: float
    end_strength: float
    end_mask_strength: float
    use_rife: bool
    rife_frames: int
    use_crossfade: bool
    text_encoder_quantization: str
    global_resident_models: bool
    vae_tiling: bool
    gpu_size: str
    gpu_duration: int
    concurrency_limit: int

    @property
    def frame_count(self) -> int:
        requested = max(int(round(self.seconds * self.fps)), 1)
        return requested + ((4 - ((requested - 1) % 4)) % 4)


def load_space_config(path: Path | None = None) -> SpaceConfig:
    parser = configparser.ConfigParser()
    config_path = path or ROOT / "config.ini"
    if not parser.read(config_path):
        raise FileNotFoundError(f"Configuration not found: {config_path}")
    section = parser["SPACE"]
    gpu_size = os.environ.get("ZERO_GPU_SIZE", section.get("zero_gpu_size", "large")).strip().lower()
    if gpu_size not in {"large", "xlarge"}:
        raise ValueError("ZERO_GPU_SIZE must be 'large' or 'xlarge'.")
    width, height = section.getint("width"), section.getint("height")
    if width % 16 or height % 16:
        raise ValueError("width and height must be divisible by 16.")
    steps, split = section.getint("steps"), section.getint("split_step")
    if not 0 < split < steps:
        raise ValueError("split_step must be between 1 and steps-1.")
    text_encoder_quantization = section.get("text_encoder_quantization", "int8_weight_only").strip().lower()
    if text_encoder_quantization not in {"int8_weight_only", "bf16"}:
        raise ValueError("text_encoder_quantization must be 'int8_weight_only' or 'bf16'.")
    return SpaceConfig(
        width=width,
        height=height,
        seconds=section.getfloat("seconds"),
        fps=section.getint("fps"),
        steps=steps,
        split_step=split,
        cfg=section.getfloat("cfgscale"),
        shift=section.getfloat("model_sampling_shift"),
        sampler=section.get("sampler_name", "euler"),
        scheduler=section.get("scheduler", "normal"),
        start_strength=section.getfloat("start_latent_strength"),
        end_strength=section.getfloat("end_latent_strength"),
        end_mask_strength=section.getfloat("end_temporal_mask_strength"),
        use_rife=_bool(section.get("apply_rife_seam", "true")),
        rife_frames=section.getint("rife_seam_frames"),
        use_crossfade=_bool(section.get("apply_crossfade", "true")),
        text_encoder_quantization=text_encoder_quantization,
        global_resident_models=_bool(section.get("global_resident_models", "false")),
        vae_tiling=_bool(section.get("vae_tiling", "false")),
        gpu_size=gpu_size,
        gpu_duration=int(os.environ.get("ZERO_GPU_DURATION", section.get("zero_gpu_duration", "300"))),
        concurrency_limit=section.getint("concurrency_limit", 1),
    )