230 lines
8.1 KiB
Python
230 lines
8.1 KiB
Python
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
|