revids/app/service.py

from __future__ import annotations

import asyncio
import logging
import os
import uuid
from collections.abc import Iterator
from datetime import datetime, timezone
from typing import Optional

import torch
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_pipelines import DistilledPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from PIL import Image

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

logger = logging.getLogger(__name__)


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
            logger.info("Loading DistilledPipeline...")
            lora_config = self._parse_lora_paths(settings.ltx_loras)
            quantization = self._build_quantization_policy()

            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=quantization,
            )
            logger.info("Pipeline loaded successfully")

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

    @staticmethod
    def _build_quantization_policy() -> object | None:
        q = settings.ltx_quantization
        if not q:
            return None
        if q == "fp8_cast":
            from ltx_core.quantization.fp8_cast import (
                TRANSFORMER_LINEAR_DOWNCAST_MAP,
                UPCAST_DURING_INFERENCE,
                fp8_cast_fuse_rule,
            )
            from ltx_core.quantization.policy import QuantizationPolicy

            return QuantizationPolicy(
                sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
                module_ops=(UPCAST_DURING_INFERENCE,),
                fuse_rule=fp8_cast_fuse_rule,
            )
        if q == "fp8_scaled_mm":
            from ltx_core.quantization.fp8_scaled_mm import (
                FP8_SCALED_MM_MODULE_OPS,
                fp8_scaled_mm_fuse_rule,
            )
            from ltx_core.quantization.policy import QuantizationPolicy

            return QuantizationPolicy(
                module_ops=FP8_SCALED_MM_MODULE_OPS, fuse_rule=fp8_scaled_mm_fuse_rule
            )
        return None

    async def submit_job(
        self, image: Image.Image, request_params: dict[str, object]
    ) -> 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": float(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[str, object],
    ) -> None:
        async with self._job_semaphore:
            try:
                logger.info(f"Job {job_id}: starting processing")
                await self.db.update_status(job_id, "processing")
                await self.load_pipeline()

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

                # Save image to temp file (ImageConditioningInput expects a path)
                if image.mode != "RGB":
                    image = image.convert("RGB")
                image = image.resize((width, height), Image.LANCZOS)

                img_path = os.path.join(
                    settings.video_output_dir, f"{job_id}_input.png"
                )
                os.makedirs(settings.video_output_dir, exist_ok=True)
                image.save(img_path)

                conditioning = [
                    ImageConditioningInput(
                        path=img_path, frame_idx=0, strength=1.0, crf=33
                    )
                ]

                logger.info(f"Job {job_id}: running inference")
                result = await asyncio.to_thread(
                    self.pipeline,
                    prompt=prompt,
                    seed=seed,
                    width=width,
                    height=height,
                    num_frames=num_frames,
                    frame_rate=float(fps),
                    images=conditioning,
                )

                # result is tuple[Iterator[torch.Tensor], Audio]
                video_frames: Iterator[torch.Tensor] = result[0]
                video_path = os.path.join(self._video_dir, f"{job_id}.mp4")
                await asyncio.to_thread(self._save_video, video_frames, video_path, fps)

                # Clean up temp input image
                if os.path.exists(img_path):
                    os.remove(img_path)

                logger.info(f"Job {job_id}: completed")
                await self.db.update_status(job_id, "completed")

            except Exception as e:
                logger.exception(f"Job {job_id}: failed")
                await self.db.update_status(job_id, "failed", error=str(e))

    @staticmethod
    def _save_video(
        frames_iter: Iterator[torch.Tensor], path: str, fps: float
    ) -> None:
        import imageio.v3 as iio

        frames = next(frames_iter).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