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 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() q = settings.ltx_gemma_quantization or None 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, gemma_quantization=q, ) 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 build_policy as fp8_build return fp8_build(settings.ltx_distilled_checkpoint) 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") 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