revids/app/service.py

from __future__ import annotations

import asyncio
import os
import uuid
from datetime import datetime, timezone
from typing import Optional

import torch
from ltx_pipelines import DistilledPipeline
from PIL import Image

from app.config import settings
from app.database import JobDB, JobRecord, VALID_FRAME_COUNTS


class LtxService:
    def __init__(self, db: JobDB) -> None:
        self.db = db
        self.pipeline: Optional[DistilledPipeline] = None
        self._load_lock = asyncio.Lock()
        self._job_semaphore = asyncio.Semaphore(settings.max_concurrent_jobs)
        self._video_dir = settings.video_output_dir

    async def load_pipeline(self) -> None:
        async with self._load_lock:
            if self.pipeline is not None:
                return
            lora_config = self._parse_lora_paths(settings.ltx_loras)
            self.pipeline = await asyncio.to_thread(
                DistilledPipeline,
                distilled_checkpoint_path=settings.ltx_distilled_checkpoint,
                gemma_root=settings.ltx_gemma_root,
                spatial_upsampler_path=settings.ltx_spatial_upsampler,
                loras=lora_config,
                device=settings.ltx_device,
                quantization=settings.ltx_quantization,
            )

    @staticmethod
    def _parse_lora_paths(lora_specs: list[str]) -> list[dict]:
        result = []
        for spec in lora_specs:
            if ":" in spec:
                path, scale = spec.rsplit(":", 1)
                result.append({"path": path.strip(), "scale": float(scale.strip())})
            else:
                result.append({"path": spec.strip(), "scale": 1.0})
        return result

    async def submit_job(self, image: Image.Image, request_params: dict) -> JobRecord:
        job_id = uuid.uuid4().hex[:12]
        now = datetime.now(timezone.utc).isoformat()
        prompt = request_params.get("prompt")
        width = request_params.get("width", settings.default_width)
        height = request_params.get("height", settings.default_height)
        num_frames = request_params.get("num_frames", settings.default_frames)
        fps = request_params.get("fps", settings.default_fps)
        seed = request_params.get("seed")

        if num_frames not in VALID_FRAME_COUNTS:
            candidates = sorted(
                f for f in VALID_FRAME_COUNTS if abs(f - num_frames) <= 8
            )
            raise ValueError(
                f"num_frames must be 8n+1. Got {num_frames}. "
                f"Closest valid: {candidates[:4]}"
            )

        job = JobRecord(
            id=job_id,
            status="pending",
            prompt=prompt,
            params={
                "width": width,
                "height": height,
                "num_frames": num_frames,
                "fps": fps,
                "seed": seed,
            },
            created_at=now,
        )
        await self.db.create_job(job)
        asyncio.create_task(self._process_job(job_id, image, job.params))
        return job

    async def _process_job(
        self,
        job_id: str,
        image: Image.Image,
        params: dict,
    ) -> None:
        async with self._job_semaphore:
            try:
                await self.db.update_status(job_id, "processing")
                await self.load_pipeline()

                prompt = params.get("prompt", "") or ""
                width = params.get("width", settings.default_width)
                height = params.get("height", settings.default_height)
                num_frames = params.get("num_frames", settings.default_frames)
                fps = params.get("fps", settings.default_fps)
                seed = params.get("seed") if params.get("seed") is not None else 42

                if image.mode != "RGB":
                    image = image.convert("RGB")
                image = image.resize((width, height), Image.LANCZOS)

                output = await asyncio.to_thread(
                    self.pipeline,
                    prompt=prompt,
                    images=image,
                    width=width,
                    height=height,
                    num_frames=num_frames,
                    frame_rate=fps,
                    seed=seed,
                )

                video_path = os.path.join(self._video_dir, f"{job_id}.mp4")
                os.makedirs(self._video_dir, exist_ok=True)
                await asyncio.to_thread(self._save_video, output, video_path, fps)
                await self.db.update_status(job_id, "completed")

            except Exception as e:
                await self.db.update_status(job_id, "failed", error=str(e))

    @staticmethod
    def _save_video(output: any, path: str, fps: float) -> None:
        import imageio.v3 as iio

        if isinstance(output, torch.Tensor):
            frames = output.cpu().numpy()
        elif isinstance(output, dict) and "video" in output:
            v = output["video"]
            frames = v.cpu().numpy() if isinstance(v, torch.Tensor) else v
        else:
            frames = output
            if isinstance(frames, torch.Tensor):
                frames = frames.cpu().numpy()

        if frames.ndim == 4:
            frames = frames[0]
        if frames.max() <= 1.0:
            frames = (frames * 255).astype("uint8")
        elif frames.dtype != "uint8":
            frames = frames.clip(0, 255).astype("uint8")

        iio.imwrite(path, frames, fps=float(fps), codec="libx264")

    async def get_job(self, job_id: str) -> JobRecord | None:
        return await self.db.get_job(job_id)

    async def delete_job(self, job_id: str) -> tuple[bool, str | None]:
        video_path = os.path.join(self._video_dir, f"{job_id}.mp4")
        deleted = await self.db.delete_job(job_id)
        if os.path.exists(video_path):
            os.remove(video_path)
        return deleted, job_id

    async def get_video_path(self, job_id: str) -> str | None:
        video_path = os.path.join(self._video_dir, f"{job_id}.mp4")
        if os.path.exists(video_path):
            return video_path
        return None


_service: Optional[LtxService] = None


async def get_service() -> LtxService:
    global _service
    if _service is None:
        db = JobDB(settings.db_path)
        await db.init()
        _service = LtxService(db)
    return _service