wehub-resource-sync
905 lines
31 KiB
Python
Executable File
905 lines
31 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import argparse
|
|
import datetime
|
|
import gc
|
|
import math
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import traceback
|
|
from concurrent import futures
|
|
|
|
import grpc
|
|
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
|
|
from longcat_utils import (
|
|
BASE_MODEL_ID,
|
|
MODEL_KIND_AVATAR,
|
|
MODEL_KIND_BASE,
|
|
attention_overrides,
|
|
avatar_segments_for_duration,
|
|
avatar_segments_for_frames,
|
|
classify_model,
|
|
normalize_model_source,
|
|
normalize_num_frames,
|
|
parse_options,
|
|
require_bool,
|
|
require_float,
|
|
require_int,
|
|
validate_dimensions,
|
|
)
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "common"))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "common"))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "sources", "LongCat-Video"))
|
|
|
|
from grpc_auth import get_auth_interceptors
|
|
|
|
|
|
MAX_WORKERS = int(os.environ.get("PYTHON_GRPC_MAX_WORKERS", "1"))
|
|
|
|
DEFAULT_NEGATIVE_PROMPT = (
|
|
"Close-up, bright tones, overexposed, static, blurred details, subtitles, "
|
|
"paintings, low quality, JPEG compression residue, ugly, incomplete, extra "
|
|
"fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
|
|
"misshapen limbs, fused fingers, still picture, messy background, three legs, "
|
|
"many people in the background, walking backwards"
|
|
)
|
|
|
|
LOAD_OPTIONS = {
|
|
"attention_backend",
|
|
"base_model",
|
|
"max_segments",
|
|
"resolution",
|
|
"use_distill",
|
|
"use_int8",
|
|
}
|
|
|
|
REQUEST_PARAMS = {
|
|
"audio_guidance_scale",
|
|
"mask_frame_range",
|
|
"num_segments",
|
|
"offload_kv_cache",
|
|
"ref_img_index",
|
|
"resolution",
|
|
}
|
|
|
|
BASE_CHECKPOINT_PATTERNS = [
|
|
"config.json",
|
|
"model_index.json",
|
|
"dit/**",
|
|
"lora/cfg_step_lora.safetensors",
|
|
"scheduler/**",
|
|
"text_encoder/**",
|
|
"tokenizer/**",
|
|
"vae/**",
|
|
]
|
|
|
|
AVATAR_BASE_PATTERNS = [
|
|
"config.json",
|
|
"model_index.json",
|
|
"text_encoder/**",
|
|
"tokenizer/**",
|
|
"vae/**",
|
|
]
|
|
|
|
AVATAR_COMMON_PATTERNS = [
|
|
"config.json",
|
|
"lora/dmd_lora.safetensors",
|
|
"model_index.json",
|
|
"scheduler/**",
|
|
"whisper-large-v3/config.json",
|
|
"whisper-large-v3/model.safetensors",
|
|
"whisper-large-v3/preprocessor_config.json",
|
|
]
|
|
|
|
|
|
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
|
def __init__(self):
|
|
self.model_kind = None
|
|
self.pipeline = None
|
|
self.options = {}
|
|
self.device_index = 0
|
|
self.cp_split_hw = None
|
|
self._dist_store_dir = None
|
|
|
|
def Health(self, request, context):
|
|
return backend_pb2.Reply(message=b"OK")
|
|
|
|
def LoadModel(self, request, context):
|
|
model = request.Model
|
|
if request.ModelFile and os.path.isdir(request.ModelFile):
|
|
model = request.ModelFile
|
|
|
|
model_kind = classify_model(model)
|
|
if model_kind is None:
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INVALID_ARGUMENT,
|
|
"longcat-video only accepts LongCat-Video or LongCat-Video-Avatar-1.5 checkpoints",
|
|
)
|
|
|
|
try:
|
|
options = parse_options(request.Options)
|
|
unknown = sorted(set(options) - LOAD_OPTIONS)
|
|
if unknown:
|
|
raise ValueError(f"unknown model option(s): {', '.join(unknown)}")
|
|
|
|
self._import_torch()
|
|
if not self.torch.cuda.is_available():
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.FAILED_PRECONDITION,
|
|
"longcat-video requires an NVIDIA CUDA GPU",
|
|
)
|
|
if request.TensorParallelSize > 1:
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.UNIMPLEMENTED,
|
|
"longcat-video currently supports one GPU per backend process",
|
|
)
|
|
self._import_runtime()
|
|
|
|
attention_name = str(options.get("attention_backend", "sdpa")).lower()
|
|
attention_overrides(attention_name)
|
|
resolution = str(options.get("resolution", "480p")).lower()
|
|
if resolution not in {"480p", "720p"}:
|
|
raise ValueError("resolution must be 480p or 720p")
|
|
|
|
use_distill_default = model_kind == MODEL_KIND_AVATAR
|
|
use_distill = require_bool(
|
|
options.get("use_distill", use_distill_default),
|
|
"use_distill",
|
|
)
|
|
use_int8 = require_bool(options.get("use_int8", False), "use_int8")
|
|
if model_kind == MODEL_KIND_BASE and use_int8:
|
|
raise ValueError(
|
|
"use_int8 is supported only by LongCat-Video-Avatar-1.5"
|
|
)
|
|
|
|
self.options = {
|
|
**options,
|
|
"attention_backend": attention_name,
|
|
"resolution": resolution,
|
|
"use_distill": use_distill,
|
|
"use_int8": use_int8,
|
|
"max_segments": require_int(
|
|
options.get("max_segments", 8),
|
|
"max_segments",
|
|
minimum=1,
|
|
maximum=64,
|
|
),
|
|
}
|
|
|
|
self._release_model()
|
|
self._ensure_distributed()
|
|
if model_kind == MODEL_KIND_BASE:
|
|
self._load_base_model(model)
|
|
else:
|
|
self._load_avatar_model(model)
|
|
self.model_kind = model_kind
|
|
print(
|
|
f"Loaded {normalize_model_source(model)} as {model_kind} "
|
|
f"with attention_backend={attention_name}",
|
|
file=sys.stderr,
|
|
)
|
|
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
|
except ValueError as err:
|
|
self._release_model()
|
|
return self._fail(context, grpc.StatusCode.INVALID_ARGUMENT, str(err))
|
|
except Exception as err:
|
|
self._release_model()
|
|
print(f"Error loading LongCat model: {err}", file=sys.stderr)
|
|
traceback.print_exc()
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INTERNAL,
|
|
f"failed to load LongCat model: {err}",
|
|
)
|
|
|
|
def Free(self, request, context):
|
|
self._release_model()
|
|
return backend_pb2.Result(message="Model released", success=True)
|
|
|
|
def GenerateVideo(self, request, context):
|
|
if self.pipeline is None or self.model_kind is None:
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.FAILED_PRECONDITION,
|
|
"model is not loaded",
|
|
)
|
|
if not request.prompt.strip():
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INVALID_ARGUMENT,
|
|
"prompt is required",
|
|
)
|
|
if not request.dst:
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INVALID_ARGUMENT,
|
|
"output destination is required",
|
|
)
|
|
if request.end_image:
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INVALID_ARGUMENT,
|
|
"longcat-video does not support end_image conditioning",
|
|
)
|
|
|
|
request_state = {"finished": False}
|
|
|
|
def interrupt_if_cancelled():
|
|
if not request_state["finished"] and self.pipeline is not None:
|
|
self.pipeline._interrupt = True
|
|
|
|
try:
|
|
params = dict(request.params)
|
|
unknown = sorted(set(params) - REQUEST_PARAMS)
|
|
if unknown:
|
|
raise ValueError(f"unknown request param(s): {', '.join(unknown)}")
|
|
|
|
os.makedirs(os.path.dirname(request.dst) or ".", mode=0o750, exist_ok=True)
|
|
if hasattr(context, "add_callback"):
|
|
context.add_callback(interrupt_if_cancelled)
|
|
|
|
if request.start_image and not os.path.isfile(request.start_image):
|
|
raise ValueError("start_image is not a readable staged file")
|
|
if request.num_frames < 0:
|
|
raise ValueError("num_frames must not be negative")
|
|
|
|
if self.model_kind == MODEL_KIND_BASE:
|
|
if request.audio:
|
|
raise ValueError(
|
|
"audio input requires a LongCat-Video-Avatar-1.5 model"
|
|
)
|
|
self._generate_base(request, params)
|
|
else:
|
|
self._generate_avatar(request, params, context)
|
|
|
|
return backend_pb2.Result(
|
|
message="Video generated successfully", success=True
|
|
)
|
|
except ValueError as err:
|
|
return self._fail(context, grpc.StatusCode.INVALID_ARGUMENT, str(err))
|
|
except Exception as err:
|
|
print(f"Error generating LongCat video: {err}", file=sys.stderr)
|
|
traceback.print_exc()
|
|
return self._fail(
|
|
context,
|
|
grpc.StatusCode.INTERNAL,
|
|
f"LongCat video generation failed: {err}",
|
|
)
|
|
finally:
|
|
request_state["finished"] = True
|
|
if self.pipeline is not None:
|
|
self.pipeline._interrupt = False
|
|
|
|
def _import_torch(self):
|
|
if hasattr(self, "torch"):
|
|
return
|
|
|
|
import torch
|
|
|
|
self.torch = torch
|
|
|
|
def _import_runtime(self):
|
|
if hasattr(self, "LongCatVideoPipeline"):
|
|
return
|
|
|
|
import imageio.v2 as imageio
|
|
import imageio_ffmpeg
|
|
import librosa
|
|
import numpy as np
|
|
import torch.distributed as dist
|
|
from diffusers.utils import load_image
|
|
from huggingface_hub import snapshot_download
|
|
from PIL import Image
|
|
from transformers import AutoTokenizer, UMT5EncoderModel
|
|
|
|
from longcat_video.audio_process import (
|
|
get_audio_encoder,
|
|
get_audio_feature_extractor,
|
|
)
|
|
from longcat_video.context_parallel import context_parallel_util
|
|
from longcat_video.modules.autoencoder_kl_wan import AutoencoderKLWan
|
|
from longcat_video.modules.avatar.longcat_video_dit_avatar import (
|
|
LongCatVideoAvatarTransformer3DModel,
|
|
)
|
|
from longcat_video.modules.longcat_video_dit import (
|
|
LongCatVideoTransformer3DModel,
|
|
)
|
|
from longcat_video.modules.quantization import load_quantized_dit
|
|
from longcat_video.modules.scheduling_flow_match_euler_discrete import (
|
|
FlowMatchEulerDiscreteScheduler,
|
|
)
|
|
from longcat_video.pipeline_longcat_video import LongCatVideoPipeline
|
|
from longcat_video.pipeline_longcat_video_avatar import (
|
|
LongCatVideoAvatarPipeline,
|
|
)
|
|
|
|
self.imageio = imageio
|
|
self.imageio_ffmpeg = imageio_ffmpeg
|
|
self.librosa = librosa
|
|
self.np = np
|
|
self.dist = dist
|
|
self.load_image = load_image
|
|
self.snapshot_download = snapshot_download
|
|
self.Image = Image
|
|
self.AutoTokenizer = AutoTokenizer
|
|
self.UMT5EncoderModel = UMT5EncoderModel
|
|
self.get_audio_encoder = get_audio_encoder
|
|
self.get_audio_feature_extractor = get_audio_feature_extractor
|
|
self.context_parallel_util = context_parallel_util
|
|
self.AutoencoderKLWan = AutoencoderKLWan
|
|
self.LongCatVideoAvatarTransformer3DModel = LongCatVideoAvatarTransformer3DModel
|
|
self.LongCatVideoTransformer3DModel = LongCatVideoTransformer3DModel
|
|
self.load_quantized_dit = load_quantized_dit
|
|
self.FlowMatchEulerDiscreteScheduler = FlowMatchEulerDiscreteScheduler
|
|
self.LongCatVideoPipeline = LongCatVideoPipeline
|
|
self.LongCatVideoAvatarPipeline = LongCatVideoAvatarPipeline
|
|
|
|
def _ensure_distributed(self):
|
|
self.torch.cuda.set_device(self.device_index)
|
|
if not self.dist.is_initialized():
|
|
self._dist_store_dir = tempfile.mkdtemp(prefix="localai-longcat-dist-")
|
|
init_file = os.path.join(self._dist_store_dir, "store")
|
|
self.dist.init_process_group(
|
|
backend="nccl",
|
|
init_method=f"file://{init_file}",
|
|
rank=0,
|
|
world_size=1,
|
|
timeout=datetime.timedelta(hours=24),
|
|
)
|
|
self.context_parallel_util.init_context_parallel(
|
|
context_parallel_size=1,
|
|
global_rank=0,
|
|
world_size=1,
|
|
)
|
|
self.cp_split_hw = self.context_parallel_util.get_optimal_split(1)
|
|
|
|
def _resolve_checkpoint(self, model, patterns):
|
|
source = normalize_model_source(model)
|
|
if os.path.isdir(source):
|
|
return source
|
|
print(f"Downloading required files for {source}", file=sys.stderr)
|
|
return self.snapshot_download(repo_id=source, allow_patterns=patterns)
|
|
|
|
def _load_base_model(self, model):
|
|
checkpoint = self._resolve_checkpoint(model, BASE_CHECKPOINT_PATTERNS)
|
|
dtype = self.torch.bfloat16
|
|
overrides = attention_overrides(self.options["attention_backend"])
|
|
|
|
tokenizer = self.AutoTokenizer.from_pretrained(
|
|
checkpoint,
|
|
subfolder="tokenizer",
|
|
)
|
|
text_encoder = self.UMT5EncoderModel.from_pretrained(
|
|
checkpoint,
|
|
subfolder="text_encoder",
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
vae = self.AutoencoderKLWan.from_pretrained(
|
|
checkpoint,
|
|
subfolder="vae",
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
scheduler = self.FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
checkpoint,
|
|
subfolder="scheduler",
|
|
)
|
|
dit = self.LongCatVideoTransformer3DModel.from_pretrained(
|
|
checkpoint,
|
|
subfolder="dit",
|
|
cp_split_hw=self.cp_split_hw,
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
**overrides,
|
|
)
|
|
if self.options["use_distill"]:
|
|
dit.load_lora(
|
|
os.path.join(checkpoint, "lora", "cfg_step_lora.safetensors"),
|
|
"cfg_step_lora",
|
|
)
|
|
dit.enable_loras(["cfg_step_lora"])
|
|
|
|
self.pipeline = self.LongCatVideoPipeline(
|
|
tokenizer=tokenizer,
|
|
text_encoder=text_encoder,
|
|
vae=vae,
|
|
scheduler=scheduler,
|
|
dit=dit,
|
|
)
|
|
self.pipeline.to(self.device_index)
|
|
|
|
def _load_avatar_model(self, model):
|
|
avatar_patterns = list(AVATAR_COMMON_PATTERNS)
|
|
model_subfolder = (
|
|
"base_model_int8" if self.options["use_int8"] else "base_model"
|
|
)
|
|
avatar_patterns.append(f"{model_subfolder}/**")
|
|
checkpoint = self._resolve_checkpoint(model, avatar_patterns)
|
|
|
|
base_model = self.options.get("base_model")
|
|
if not base_model and os.path.isdir(normalize_model_source(model)):
|
|
sibling = os.path.join(
|
|
os.path.dirname(normalize_model_source(model)), "LongCat-Video"
|
|
)
|
|
if os.path.isdir(sibling):
|
|
base_model = sibling
|
|
base_model = base_model or BASE_MODEL_ID
|
|
if classify_model(str(base_model)) != MODEL_KIND_BASE:
|
|
raise ValueError("base_model must point to a LongCat-Video checkpoint")
|
|
base_checkpoint = self._resolve_checkpoint(base_model, AVATAR_BASE_PATTERNS)
|
|
|
|
dtype = self.torch.bfloat16
|
|
overrides = attention_overrides(self.options["attention_backend"])
|
|
tokenizer = self.AutoTokenizer.from_pretrained(
|
|
base_checkpoint,
|
|
subfolder="tokenizer",
|
|
)
|
|
text_encoder = self.UMT5EncoderModel.from_pretrained(
|
|
base_checkpoint,
|
|
subfolder="text_encoder",
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
vae = self.AutoencoderKLWan.from_pretrained(
|
|
base_checkpoint,
|
|
subfolder="vae",
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
scheduler = self.FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
checkpoint,
|
|
subfolder="scheduler",
|
|
)
|
|
|
|
if self.options["use_int8"]:
|
|
previous_dtype = self.torch.get_default_dtype()
|
|
self.torch.set_default_dtype(dtype)
|
|
try:
|
|
dit = self.load_quantized_dit(
|
|
checkpoint,
|
|
subfolder="base_model_int8",
|
|
cp_split_hw=self.cp_split_hw,
|
|
**overrides,
|
|
)
|
|
finally:
|
|
self.torch.set_default_dtype(previous_dtype)
|
|
else:
|
|
dit = self.LongCatVideoAvatarTransformer3DModel.from_pretrained(
|
|
checkpoint,
|
|
subfolder="base_model",
|
|
cp_split_hw=self.cp_split_hw,
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
**overrides,
|
|
)
|
|
|
|
if self.options["use_distill"]:
|
|
dit.load_lora(
|
|
os.path.join(checkpoint, "lora", "dmd_lora.safetensors"),
|
|
"dmd",
|
|
multiplier=1.0,
|
|
lora_network_dim=128,
|
|
lora_network_alpha=64,
|
|
)
|
|
dit.enable_loras(["dmd"])
|
|
|
|
audio_checkpoint = os.path.join(checkpoint, "whisper-large-v3")
|
|
audio_encoder = self.get_audio_encoder(
|
|
audio_checkpoint,
|
|
MODEL_KIND_AVATAR + "-v1.5",
|
|
).to(self.device_index)
|
|
audio_feature_extractor = self.get_audio_feature_extractor(
|
|
audio_checkpoint,
|
|
MODEL_KIND_AVATAR + "-v1.5",
|
|
)
|
|
self.pipeline = self.LongCatVideoAvatarPipeline(
|
|
tokenizer=tokenizer,
|
|
text_encoder=text_encoder,
|
|
vae=vae,
|
|
scheduler=scheduler,
|
|
dit=dit,
|
|
audio_encoder=audio_encoder,
|
|
audio_feature_extractor=audio_feature_extractor,
|
|
model_type="avatar-v1.5",
|
|
)
|
|
self.pipeline.to(self.device_index)
|
|
|
|
def _generate_base(self, request, params):
|
|
use_distill = self.options["use_distill"]
|
|
frames = normalize_num_frames(request.num_frames)
|
|
steps = (
|
|
16
|
|
if use_distill
|
|
else require_int(
|
|
request.step or 50,
|
|
"step",
|
|
minimum=1,
|
|
maximum=200,
|
|
)
|
|
)
|
|
guidance_scale = (
|
|
1.0
|
|
if use_distill
|
|
else require_float(
|
|
request.cfg_scale or 4.0,
|
|
"cfg_scale",
|
|
minimum=0.0,
|
|
maximum=30.0,
|
|
)
|
|
)
|
|
fps = require_int(request.fps or 15, "fps", minimum=1, maximum=60)
|
|
seed = request.seed if request.seed > 0 else 42
|
|
negative_prompt = request.negative_prompt or DEFAULT_NEGATIVE_PROMPT
|
|
generator = self.torch.Generator(device=self.device_index).manual_seed(seed)
|
|
|
|
if request.start_image:
|
|
resolution = self._resolution(params)
|
|
image = self.load_image(request.start_image)
|
|
output = self.pipeline.generate_i2v(
|
|
image=image,
|
|
prompt=request.prompt,
|
|
negative_prompt=negative_prompt,
|
|
resolution=resolution,
|
|
num_frames=frames,
|
|
num_inference_steps=steps,
|
|
use_distill=use_distill,
|
|
guidance_scale=guidance_scale,
|
|
generator=generator,
|
|
)[0]
|
|
else:
|
|
width, height = validate_dimensions(request.width, request.height)
|
|
output = self.pipeline.generate_t2v(
|
|
prompt=request.prompt,
|
|
negative_prompt=negative_prompt,
|
|
height=height,
|
|
width=width,
|
|
num_frames=frames,
|
|
num_inference_steps=steps,
|
|
use_distill=use_distill,
|
|
guidance_scale=guidance_scale,
|
|
generator=generator,
|
|
)[0]
|
|
|
|
self._save_video(output, request.dst, fps)
|
|
|
|
def _generate_avatar(self, request, params, context):
|
|
if not request.audio:
|
|
raise ValueError("audio is required for LongCat-Video-Avatar-1.5")
|
|
if not os.path.isfile(request.audio):
|
|
raise ValueError("audio input is not a readable staged file")
|
|
|
|
use_distill = self.options["use_distill"]
|
|
steps = (
|
|
8
|
|
if use_distill
|
|
else require_int(
|
|
request.step or 50,
|
|
"step",
|
|
minimum=1,
|
|
maximum=200,
|
|
)
|
|
)
|
|
text_guidance = (
|
|
1.0
|
|
if use_distill
|
|
else require_float(
|
|
request.cfg_scale or 4.0,
|
|
"cfg_scale",
|
|
minimum=0.0,
|
|
maximum=30.0,
|
|
)
|
|
)
|
|
audio_guidance = (
|
|
1.0
|
|
if use_distill
|
|
else require_float(
|
|
params.get("audio_guidance_scale", 4.0),
|
|
"audio_guidance_scale",
|
|
minimum=0.0,
|
|
maximum=20.0,
|
|
)
|
|
)
|
|
seed = request.seed if request.seed > 0 else 42
|
|
generator = self.torch.Generator(device=self.device_index).manual_seed(seed)
|
|
negative_prompt = request.negative_prompt or DEFAULT_NEGATIVE_PROMPT
|
|
resolution = self._resolution(params)
|
|
|
|
speech, sample_rate = self.librosa.load(request.audio, sr=16000, mono=True)
|
|
if speech.size == 0:
|
|
raise ValueError("audio contains no samples")
|
|
audio_duration = len(speech) / sample_rate
|
|
segments = self._avatar_segments(request, params, audio_duration)
|
|
|
|
segment_frames = 93
|
|
conditioning_frames = 13
|
|
avatar_fps = 25
|
|
generated_duration = (
|
|
segment_frames + (segments - 1) * (segment_frames - conditioning_frames)
|
|
) / avatar_fps
|
|
pad_samples = max(
|
|
0, math.ceil((generated_duration - audio_duration) * sample_rate)
|
|
)
|
|
if pad_samples:
|
|
speech = self.np.pad(speech, (0, pad_samples))
|
|
|
|
full_audio_embedding = self.pipeline.get_audio_embedding(
|
|
speech,
|
|
fps=avatar_fps,
|
|
device=self.device_index,
|
|
sample_rate=sample_rate,
|
|
model_type="avatar-v1.5",
|
|
)
|
|
if not self.torch.isfinite(full_audio_embedding).all():
|
|
raise ValueError("audio encoder returned non-finite values")
|
|
|
|
indices = self.torch.arange(5) - 2
|
|
|
|
def audio_window(start_index):
|
|
centers = self.torch.arange(
|
|
start_index,
|
|
start_index + segment_frames,
|
|
).unsqueeze(1) + indices.unsqueeze(0)
|
|
centers = self.torch.clamp(
|
|
centers,
|
|
min=0,
|
|
max=full_audio_embedding.shape[0] - 1,
|
|
)
|
|
return full_audio_embedding[centers][None, ...].to(self.device_index)
|
|
|
|
audio_start = 0
|
|
common = {
|
|
"prompt": request.prompt,
|
|
"negative_prompt": negative_prompt,
|
|
"num_frames": segment_frames,
|
|
"num_inference_steps": steps,
|
|
"text_guidance_scale": text_guidance,
|
|
"audio_guidance_scale": audio_guidance,
|
|
"output_type": "both",
|
|
"generator": generator,
|
|
"audio_emb": audio_window(audio_start),
|
|
"use_distill": use_distill,
|
|
}
|
|
|
|
if request.start_image:
|
|
output, latent = self.pipeline.generate_ai2v(
|
|
image=self.load_image(request.start_image),
|
|
resolution=resolution,
|
|
**common,
|
|
)
|
|
else:
|
|
width, height = validate_dimensions(request.width, request.height)
|
|
output, latent = self.pipeline.generate_at2v(
|
|
height=height,
|
|
width=width,
|
|
**common,
|
|
)
|
|
|
|
video = self._frames_to_pil(output[0])
|
|
width, height = video[0].size
|
|
current_video = video
|
|
reference_latent = latent[:, :, :1].clone()
|
|
all_frames = list(video)
|
|
|
|
for segment in range(1, segments):
|
|
if hasattr(context, "is_active") and not context.is_active():
|
|
raise RuntimeError("request was cancelled")
|
|
print(
|
|
f"Generating avatar segment {segment + 1}/{segments}", file=sys.stderr
|
|
)
|
|
audio_start += segment_frames - conditioning_frames
|
|
output, latent = self.pipeline.generate_avc(
|
|
video=current_video,
|
|
video_latent=latent,
|
|
prompt=request.prompt,
|
|
negative_prompt=negative_prompt,
|
|
height=height,
|
|
width=width,
|
|
num_frames=segment_frames,
|
|
num_cond_frames=conditioning_frames,
|
|
num_inference_steps=steps,
|
|
text_guidance_scale=text_guidance,
|
|
audio_guidance_scale=audio_guidance,
|
|
generator=generator,
|
|
output_type="both",
|
|
use_kv_cache=True,
|
|
offload_kv_cache=require_bool(
|
|
params.get("offload_kv_cache", False),
|
|
"offload_kv_cache",
|
|
),
|
|
enhance_hf=not use_distill,
|
|
audio_emb=audio_window(audio_start),
|
|
ref_latent=reference_latent,
|
|
ref_img_index=require_int(
|
|
params.get("ref_img_index", 10),
|
|
"ref_img_index",
|
|
minimum=-30,
|
|
maximum=30,
|
|
),
|
|
mask_frame_range=require_int(
|
|
params.get("mask_frame_range", 3),
|
|
"mask_frame_range",
|
|
minimum=0,
|
|
maximum=32,
|
|
),
|
|
use_distill=use_distill,
|
|
)
|
|
current_video = self._frames_to_pil(output[0])
|
|
all_frames.extend(current_video[conditioning_frames:])
|
|
|
|
self._save_avatar_video(all_frames, request.audio, request.dst, avatar_fps)
|
|
|
|
def _avatar_segments(self, request, params, audio_duration):
|
|
if "num_segments" in params:
|
|
segments = require_int(
|
|
params["num_segments"],
|
|
"num_segments",
|
|
minimum=1,
|
|
)
|
|
elif request.num_frames > 0:
|
|
segments = avatar_segments_for_frames(request.num_frames)
|
|
else:
|
|
segments = avatar_segments_for_duration(audio_duration)
|
|
|
|
max_segments = self.options["max_segments"]
|
|
if segments > max_segments:
|
|
raise ValueError(
|
|
f"request needs {segments} avatar segments, but max_segments is {max_segments}; "
|
|
"trim the audio or raise the model's max_segments option"
|
|
)
|
|
return segments
|
|
|
|
def _resolution(self, params):
|
|
resolution = str(params.get("resolution", self.options["resolution"])).lower()
|
|
if resolution not in {"480p", "720p"}:
|
|
raise ValueError("resolution must be 480p or 720p")
|
|
return resolution
|
|
|
|
def _frames_to_pil(self, frames):
|
|
images = []
|
|
for frame in frames:
|
|
array = self.np.asarray(frame)
|
|
if self.np.issubdtype(array.dtype, self.np.floating):
|
|
array = self.np.clip(array, 0.0, 1.0) * 255
|
|
images.append(self.Image.fromarray(array.astype(self.np.uint8)))
|
|
return images
|
|
|
|
def _save_video(self, frames, path, fps):
|
|
writer = self.imageio.get_writer(
|
|
path,
|
|
format="FFMPEG",
|
|
mode="I",
|
|
fps=fps,
|
|
codec="libx264",
|
|
macro_block_size=1,
|
|
ffmpeg_params=[
|
|
"-crf",
|
|
"18",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-movflags",
|
|
"+faststart",
|
|
"-f",
|
|
"mp4",
|
|
],
|
|
)
|
|
try:
|
|
for frame in frames:
|
|
array = self.np.asarray(frame)
|
|
if self.np.issubdtype(array.dtype, self.np.floating):
|
|
array = self.np.clip(array, 0.0, 1.0) * 255
|
|
writer.append_data(array.astype(self.np.uint8))
|
|
finally:
|
|
writer.close()
|
|
|
|
def _save_avatar_video(self, frames, audio_path, dst, fps):
|
|
output_dir = os.path.dirname(dst) or "."
|
|
handle, silent_path = tempfile.mkstemp(
|
|
prefix="longcat-silent-",
|
|
suffix=".mp4",
|
|
dir=output_dir,
|
|
)
|
|
os.close(handle)
|
|
try:
|
|
self._save_video(frames, silent_path, fps)
|
|
command = [
|
|
self.imageio_ffmpeg.get_ffmpeg_exe(),
|
|
"-y",
|
|
"-i",
|
|
silent_path,
|
|
"-i",
|
|
audio_path,
|
|
"-map",
|
|
"0:v:0",
|
|
"-map",
|
|
"1:a:0",
|
|
"-c:v",
|
|
"copy",
|
|
"-c:a",
|
|
"aac",
|
|
"-b:a",
|
|
"192k",
|
|
"-shortest",
|
|
"-movflags",
|
|
"+faststart",
|
|
"-f",
|
|
"mp4",
|
|
dst,
|
|
]
|
|
subprocess.run(
|
|
command,
|
|
check=True,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
except subprocess.CalledProcessError as err:
|
|
details = (err.stderr or "ffmpeg failed")[-2000:]
|
|
raise RuntimeError(f"failed to mux avatar audio: {details}") from err
|
|
finally:
|
|
try:
|
|
os.remove(silent_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
def _release_model(self):
|
|
self.pipeline = None
|
|
self.model_kind = None
|
|
gc.collect()
|
|
if hasattr(self, "torch") and self.torch.cuda.is_available():
|
|
self.torch.cuda.empty_cache()
|
|
self.torch.cuda.ipc_collect()
|
|
|
|
@staticmethod
|
|
def _fail(context, code, message):
|
|
if context is not None:
|
|
context.set_code(code)
|
|
context.set_details(message)
|
|
return backend_pb2.Result(message=message, success=False)
|
|
|
|
|
|
def serve(address):
|
|
server = grpc.server(
|
|
futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
|
options=[
|
|
("grpc.max_message_length", 64 * 1024 * 1024),
|
|
("grpc.max_send_message_length", 64 * 1024 * 1024),
|
|
("grpc.max_receive_message_length", 64 * 1024 * 1024),
|
|
],
|
|
interceptors=get_auth_interceptors(),
|
|
)
|
|
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
|
server.add_insecure_port(address)
|
|
server.start()
|
|
print(f"LongCat Video backend listening on {address}", file=sys.stderr)
|
|
|
|
def stop_server(signum, frame):
|
|
del signum, frame
|
|
server.stop(0)
|
|
|
|
signal.signal(signal.SIGINT, stop_server)
|
|
signal.signal(signal.SIGTERM, stop_server)
|
|
server.wait_for_termination()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run the LongCat Video gRPC backend")
|
|
parser.add_argument(
|
|
"--addr",
|
|
default="localhost:50051",
|
|
help="address on which to serve the backend",
|
|
)
|
|
arguments = parser.parse_args()
|
|
serve(arguments.addr)
|