| import json |
| import math |
| import os |
| from typing import Optional, List |
|
|
| import numpy as np |
| import torch |
| from accelerate import init_empty_weights |
| from mmgp import offload |
| from tqdm import tqdm |
| import librosa |
| import pyloudnorm as pyln |
| import scipy.signal as ss |
| import torch.nn.functional as F |
| from transformers import AutoFeatureExtractor, Wav2Vec2FeatureExtractor, WhisperModel |
|
|
| from shared.utils import files_locator as fl |
| from ..wan.modules.t5 import T5EncoderModel |
| from .modules.longcat_video_dit import LongCatVideoTransformer3DModel |
| from .modules.avatar.longcat_video_dit_avatar import LongCatVideoAvatarTransformer3DModel |
| from .modules.autoencoder_kl_wan import AutoencoderKLWan |
| from .modules.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler |
| from .audio_process.wav2vec2 import Wav2Vec2ModelWrapper |
| from ..qwen.convert_diffusers_qwen_vae import convert_state_dict |
| from shared.utils.text_encoder_cache import TextEncoderCache |
|
|
|
|
| def _load_json_config(path): |
| with open(path, "r", encoding="utf-8") as f: |
| cfg = json.load(f) |
| cfg.pop("_class_name", None) |
| cfg.pop("_diffusers_version", None) |
| cfg.pop("architectures", None) |
| cfg.pop("model_max_length", None) |
| return cfg |
|
|
|
|
| def retrieve_latents( |
| encoder_output: torch.Tensor, |
| generator: Optional[torch.Generator] = None, |
| sample_mode: str = "sample", |
| ): |
| if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": |
| return encoder_output.latent_dist.sample(generator) |
| if hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": |
| return encoder_output.latent_dist.mode() |
| if hasattr(encoder_output, "latents"): |
| return encoder_output.latents |
| raise AttributeError("Could not access latents of provided encoder_output") |
|
|
|
|
| def optimized_scale(positive_flat, negative_flat): |
| dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) |
| squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 |
| return dot_product / squared_norm |
|
|
|
|
| LONGCAT_AVATAR_TYPES = {"longcat_avatar", "longcat_avatar_v1_5"} |
|
|
|
|
| class LongCatModel: |
| def __init__( |
| self, |
| checkpoint_dir, |
| model_filename=None, |
| model_type=None, |
| model_def=None, |
| base_model_type=None, |
| text_encoder_filename=None, |
| quantizeTransformer=False, |
| save_quantized=False, |
| dtype=torch.bfloat16, |
| VAE_dtype=torch.float32, |
| mixed_precision_transformer=False, |
| **kwargs, |
| ): |
| self.device = torch.device("cuda") |
| self.dtype = dtype |
| self.VAE_dtype = VAE_dtype |
| self.model_def = model_def or {} |
| self.base_model_type = base_model_type |
| self.is_avatar = base_model_type in LONGCAT_AVATAR_TYPES |
| self.is_avatar_v1_5 = base_model_type == "longcat_avatar_v1_5" |
| self.audio_encoder_name = None |
| self.sparse_attention_enabled = bool(self.model_def.get("sparse_attention", False)) |
| self._interrupt = False |
| self._reference_image = None |
|
|
| text_encoder_path = text_encoder_filename or fl.locate_file( |
| "umt5-xxl/models_t5_umt5-xxl-enc-bf16.safetensors", True |
| ) |
| text_encoder_folder = self.model_def.get("text_encoder_folder") |
| if text_encoder_folder: |
| tokenizer_path = fl.locate_folder(text_encoder_folder) |
| else: |
| tokenizer_path = os.path.dirname(text_encoder_path) |
| self.text_encoder = T5EncoderModel( |
| text_len=512, |
| dtype=dtype, |
| device=torch.device("cpu"), |
| checkpoint_path=text_encoder_path, |
| tokenizer_path=tokenizer_path, |
| ) |
| self.text_encoder_cache = TextEncoderCache() |
|
|
| transformer_config_path = self.model_def.get("transformer_config") |
| if not transformer_config_path: |
| transformer_config_path = ( |
| "models/longcat/configs/longcat_avatar.json" |
| if self.is_avatar |
| else "models/longcat/configs/longcat_video.json" |
| ) |
| transformer_cfg = _load_json_config(transformer_config_path) |
| if self.sparse_attention_enabled: |
| transformer_cfg["enable_bsa"] = True |
| sparse_params = self.model_def.get("sparse_attention_params") |
| if isinstance(sparse_params, dict) and sparse_params: |
| bsa_params = dict(transformer_cfg.get("bsa_params") or {}) |
| bsa_params.update(sparse_params) |
| transformer_cfg["bsa_params"] = bsa_params |
| transformer_cls = ( |
| LongCatVideoAvatarTransformer3DModel |
| if self.is_avatar |
| else LongCatVideoTransformer3DModel |
| ) |
| with init_empty_weights(include_buffers=True): |
| transformer = transformer_cls(**transformer_cfg) |
| model_path = model_filename[0] if isinstance(model_filename, (list, tuple)) else model_filename |
| if model_path is None: |
| raise ValueError("Missing LongCat transformer weights path.") |
| offload.load_model_data(transformer, model_path, writable_tensors=False) |
| transformer._model_dtype = dtype |
| transformer.eval().requires_grad_(False) |
| self.transformer = transformer |
| if save_quantized: |
| from wgp import save_quantized_model |
|
|
| save_quantized_model(transformer, model_type, model_path, dtype, transformer_config_path) |
|
|
| vae_cfg_path = "models/longcat/configs/longcat_vae.json" |
| vae_weights = self.model_def.get("vae_URL") |
| if vae_weights: |
| vae_weights = fl.locate_file(vae_weights) |
| else: |
| for candidate in ["Wan2.1_VAE_bf16.safetensors", "Wan2.1_VAE.safetensors", "longcat_vae_bf16.safetensors"]: |
| vae_weights = fl.locate_file(candidate, error_if_none=False) |
| if vae_weights: |
| break |
| if not vae_weights: |
| raise FileNotFoundError("Unable to locate a compatible VAE weights file for LongCat.") |
| def preprocess_vae_sd(sd): |
| return convert_state_dict(sd) |
|
|
| self.vae = offload.fast_load_transformers_model( |
| vae_weights, |
| modelClass=AutoencoderKLWan, |
| defaultConfigPath=vae_cfg_path, |
| writable_tensors=False, |
| preprocess_sd=preprocess_vae_sd, |
| default_dtype=VAE_dtype, |
| ) |
| self.vae = self.vae.to(dtype=VAE_dtype, device="cpu") |
| self.vae._model_dtype = VAE_dtype |
| self.vae._dtype = VAE_dtype |
| self.vae.eval().requires_grad_(False) |
|
|
| scheduler_cfg = _load_json_config(self.model_def.get("scheduler_config", "models/longcat/configs/longcat_scheduler.json")) |
| self.scheduler = FlowMatchEulerDiscreteScheduler(**scheduler_cfg) |
| self.num_timesteps = 1000 |
| self.num_distill_sample_steps = int(self.model_def.get("num_distill_sample_steps", 8 if self.is_avatar_v1_5 else 50)) |
|
|
| if self.is_avatar: |
| if self.is_avatar_v1_5: |
| whisper_folder_name = self.model_def.get("audio_encoder_folder", "whisper-large-v3") |
| whisper_folder = fl.locate_folder(whisper_folder_name) |
| whisper_model_path = fl.locate_file(os.path.join(whisper_folder_name, "model.safetensors")) |
| whisper_config_path = fl.locate_file(os.path.join(whisper_folder_name, "config.json")) |
| fl.locate_file(os.path.join(whisper_folder_name, "generation_config.json")) |
| fl.locate_file(os.path.join(whisper_folder_name, "preprocessor_config.json")) |
| self.audio_encoder_name = "whisper-large-v3" |
| self.audio_encoder = offload.fast_load_transformers_model( |
| whisper_model_path, |
| modelClass=WhisperModel, |
| defaultConfigPath=whisper_config_path, |
| modelPrefix="model", |
| writable_tensors=False, |
| default_dtype=dtype, |
| ignore_unused_weights=True, |
| ) |
| if hasattr(self.audio_encoder, "decoder"): |
| del self.audio_encoder.decoder |
| self.audio_encoder._model_dtype = dtype |
| self.audio_encoder.eval().requires_grad_(False) |
| self.audio_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_folder, local_files_only=True) |
| else: |
| wav2vec_folder = fl.locate_folder("chinese-wav2vec2-base") |
| self.audio_encoder_name = "wav2vec2" |
| self.audio_encoder = Wav2Vec2ModelWrapper(wav2vec_folder) |
| self.audio_encoder.eval().requires_grad_(False) |
| if hasattr(self.audio_encoder, "feature_extractor"): |
| self.audio_encoder.feature_extractor._freeze_parameters() |
| self.audio_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_folder, local_files_only=True) |
| else: |
| self.audio_encoder = None |
| self.audio_feature_extractor = None |
|
|
| self.vae_scale_factor_temporal = getattr(self.vae.config, "scale_factor_temporal", 4) |
| self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 8) |
| self.transformer._interrupt_check = lambda: self._interrupt |
|
|
| def get_loras_transformer(self, get_model_recursive_prop, **kwargs): |
| if not self.is_avatar_v1_5: |
| return [], [] |
| sample_solver = kwargs.get("sample_solver") or self.model_def.get("sample_solver", "distill") |
| lora_url = self.model_def.get("distill_lora_URL") if sample_solver == "distill" else None |
| return ([lora_url], ["1.0"]) if lora_url else ([], []) |
|
|
| def _clear_runtime_caches(self): |
| clear = getattr(self.transformer, "clear_runtime_caches", None) |
| if clear is not None: |
| clear() |
|
|
| def prepare_preview_payload(self, latents, preview_meta=None): |
| if not torch.is_tensor(latents): |
| return None |
| return {"latents": latents.float()} |
|
|
| def _apply_vae_tiling(self, VAE_tile_size): |
| if not hasattr(self.vae, "enable_tiling"): |
| return |
| if VAE_tile_size is None or VAE_tile_size == 0: |
| if hasattr(self.vae, "disable_tiling"): |
| self.vae.disable_tiling() |
| return |
| if isinstance(VAE_tile_size, dict): |
| tile = VAE_tile_size.get("tile_sample_min_size", None) |
| else: |
| tile = int(VAE_tile_size) |
| if tile and tile > 0: |
| stride = max(16, int(tile * 0.75)) |
| self.vae.enable_tiling( |
| tile_sample_min_height=tile, |
| tile_sample_min_width=tile, |
| tile_sample_stride_height=stride, |
| tile_sample_stride_width=stride, |
| ) |
|
|
| def _validate_sparse_attention(self, latents): |
| if not self.sparse_attention_enabled: |
| return |
| bsa_params = getattr(self.transformer.config, "bsa_params", None) or {} |
| chunk_shape_q = bsa_params.get("chunk_3d_shape_q") |
| chunk_shape_k = bsa_params.get("chunk_3d_shape_k") |
| chunk_shape = chunk_shape_q or chunk_shape_k |
| if not chunk_shape or latents.dim() != 5: |
| self.transformer.disable_bsa() |
| self.sparse_attention_enabled = False |
| print("Sparse attention disabled: missing BSA parameters.") |
| return |
| attn_mode = offload.shared_state.get("_attention", "auto") |
| require_grid_divisible = False |
| if attn_mode == "flash": |
| require_grid_divisible = True |
| elif attn_mode == "auto": |
| try: |
| from shared.attention import flash_attn_bsa_3d |
| except Exception: |
| flash_attn_bsa_3d = None |
| require_grid_divisible = flash_attn_bsa_3d is not None |
|
|
| patch_t, patch_h, patch_w = self.transformer.config.patch_size |
| n_t = latents.shape[2] // patch_t |
| n_h = latents.shape[3] // patch_h |
| n_w = latents.shape[4] // patch_w |
| cp_split_hw = getattr(self.transformer.config, "cp_split_hw", None) |
| if cp_split_hw: |
| if n_h % cp_split_hw[0] != 0 or n_w % cp_split_hw[1] != 0: |
| self.transformer.disable_bsa() |
| self.sparse_attention_enabled = False |
| print("Sparse attention disabled: cp_split_hw does not divide token grid.") |
| return |
| if require_grid_divisible: |
| shape_q = chunk_shape_q or chunk_shape |
| shape_k = chunk_shape_k or chunk_shape |
| if ( |
| n_t % shape_q[0] != 0 |
| or n_h % shape_q[1] != 0 |
| or n_w % shape_q[2] != 0 |
| or n_t % shape_k[0] != 0 |
| or n_h % shape_k[1] != 0 |
| or n_w % shape_k[2] != 0 |
| ): |
| self.transformer.disable_bsa() |
| self.sparse_attention_enabled = False |
| print("Sparse attention disabled: flash BSA needs token grid divisible by chunk shape.") |
|
|
| def _encode_prompt( |
| self, |
| prompt, |
| negative_prompt, |
| num_videos_per_prompt=1, |
| max_length=512, |
| device=None, |
| dtype=None, |
| ): |
| device = device or self.device |
| dtype = dtype or self.dtype |
| def encode_fn(prompts): |
| ids, mask = self.text_encoder.tokenizer( |
| prompts, |
| return_mask=True, |
| add_special_tokens=True, |
| ) |
| ids = ids.to(device) |
| mask = mask.to(device) |
| prompt_embeds = self.text_encoder.model(ids, mask).to(dtype) |
| return list(zip(prompt_embeds, mask)) |
| prompt_list = [prompt] if isinstance(prompt, str) else prompt |
| batch_size = len(prompt_list) |
| prompt_contexts = self.text_encoder_cache.encode( |
| encode_fn, |
| prompt_list, |
| device=device, |
| ) |
| prompt_embeds = torch.stack([ctx[0] for ctx in prompt_contexts], dim=0) |
| mask = torch.stack([ctx[1] for ctx in prompt_contexts], dim=0) |
| seq_len = prompt_embeds.shape[1] |
| prompt_embeds = prompt_embeds.unsqueeze(1) |
| if num_videos_per_prompt > 1: |
| prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1, 1) |
| prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, 1, seq_len, -1) |
| mask = mask.repeat(num_videos_per_prompt, 1) |
|
|
| neg_embeds = None |
| neg_mask = None |
| if negative_prompt is not None: |
| neg_list = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt |
| if len(neg_list) == 1 and batch_size > 1: |
| neg_list = neg_list * batch_size |
| neg_contexts = self.text_encoder_cache.encode( |
| encode_fn, |
| neg_list, |
| device=device, |
| ) |
| neg_embeds = torch.stack([ctx[0] for ctx in neg_contexts], dim=0) |
| neg_mask = torch.stack([ctx[1] for ctx in neg_contexts], dim=0) |
| neg_embeds = neg_embeds.unsqueeze(1) |
| if num_videos_per_prompt > 1: |
| neg_embeds = neg_embeds.repeat(1, num_videos_per_prompt, 1, 1) |
| neg_embeds = neg_embeds.view(batch_size * num_videos_per_prompt, 1, seq_len, -1) |
| neg_mask = neg_mask.repeat(num_videos_per_prompt, 1) |
|
|
| return prompt_embeds, mask, neg_embeds, neg_mask |
|
|
| def _prepare_latents( |
| self, |
| batch_size, |
| num_channels_latents, |
| height, |
| width, |
| num_frames, |
| dtype, |
| device, |
| generator, |
| latents=None, |
| image=None, |
| video=None, |
| num_cond_frames=0, |
| ): |
| num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 |
| if latents is None: |
| shape = ( |
| batch_size, |
| num_channels_latents, |
| num_latent_frames, |
| int(height) // self.vae_scale_factor_spatial, |
| int(width) // self.vae_scale_factor_spatial, |
| ) |
| latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) |
| else: |
| latents = latents.to(device=device, dtype=dtype) |
|
|
| num_cond_latents = 0 |
| if image is not None or video is not None: |
| cond_latents = [] |
| for i in range(batch_size): |
| if image is not None: |
| encoded_input = image[i].unsqueeze(0).unsqueeze(2) |
| else: |
| encoded_input = video[i][:, -num_cond_frames:].unsqueeze(0) |
| latent = retrieve_latents( |
| self.vae.encode(encoded_input), |
| generator, |
| sample_mode="argmax", |
| ) |
| cond_latents.append(latent) |
| cond_latents = torch.cat(cond_latents, dim=0).to(dtype) |
| cond_latents = self.normalize_latents(cond_latents) |
| num_cond_latents = 1 + (num_cond_frames - 1) // self.vae_scale_factor_temporal |
| latents[:, :, :num_cond_latents] = cond_latents[:, :, :num_cond_latents] |
| return latents, num_cond_latents |
|
|
| def normalize_latents(self, latents): |
| latents_mean = ( |
| torch.tensor(self.vae.config.latents_mean) |
| .view(1, self.vae.config.z_dim, 1, 1, 1) |
| .to(latents.device, latents.dtype) |
| ) |
| latents_std = ( |
| torch.tensor(self.vae.config.latents_std) |
| .view(1, self.vae.config.z_dim, 1, 1, 1) |
| .to(latents.device, latents.dtype) |
| ) |
| return (latents - latents_mean) / latents_std |
|
|
| def denormalize_latents(self, latents): |
| latents_mean = ( |
| torch.tensor(self.vae.config.latents_mean) |
| .view(1, self.vae.config.z_dim, 1, 1, 1) |
| .to(latents.device, latents.dtype) |
| ) |
| latents_std = ( |
| torch.tensor(self.vae.config.latents_std) |
| .view(1, self.vae.config.z_dim, 1, 1, 1) |
| .to(latents.device, latents.dtype) |
| ) |
| return latents * latents_std + latents_mean |
|
|
| def _loudness_norm(self, audio_array, sr=16000, lufs=-23, threshold=100): |
| meter = pyln.Meter(sr) |
| loudness = meter.integrated_loudness(audio_array) |
| if not np.isfinite(loudness) or loudness > threshold: |
| return audio_array |
| return pyln.normalize.loudness(audio_array, loudness, lufs) |
|
|
| def _add_noise_floor(self, audio_array, noise_level=0.0001): |
| noise = np.random.normal(0, noise_level, size=audio_array.shape) |
| return audio_array + noise |
|
|
| def _smooth_transients(self, audio_array, sr=16000): |
| b, a = ss.butter(3, 3000 / (sr / 2)) |
| return ss.lfilter(b, a, audio_array) |
|
|
| @staticmethod |
| def _interpolate_audio_state(audio_state, target_len): |
| if audio_state.shape[0] == target_len: |
| return audio_state |
| audio_state = audio_state.transpose(0, 1).unsqueeze(0).float() |
| audio_state = F.interpolate(audio_state, size=target_len, mode="linear", align_corners=False) |
| return audio_state.squeeze(0).transpose(0, 1) |
|
|
| @torch.no_grad() |
| def _get_audio_embedding_wav2vec(self, speech_array, fps=32, device="cpu", sample_rate=16000): |
| audio_duration = len(speech_array) / sample_rate |
| video_length = audio_duration * fps |
|
|
| speech_array = self._loudness_norm(speech_array, sample_rate) |
| speech_array = self._add_noise_floor(speech_array) |
| speech_array = self._smooth_transients(speech_array, sample_rate) |
|
|
| audio_feature = np.squeeze( |
| self.audio_feature_extractor(speech_array, sampling_rate=sample_rate).input_values |
| ) |
| audio_feature = np.nan_to_num(audio_feature, nan=0.0, posinf=0.0, neginf=0.0) |
| audio_feature = torch.from_numpy(audio_feature).float().to(device=device) |
| audio_feature = audio_feature.unsqueeze(0) |
|
|
| embeddings = self.audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True) |
| audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) |
| audio_emb = audio_emb.permute(1, 0, 2).contiguous() |
| audio_emb = torch.nan_to_num(audio_emb, nan=0.0, posinf=0.0, neginf=0.0) |
| return audio_emb |
|
|
| @torch.no_grad() |
| def _get_audio_embedding_whisper(self, speech_array, fps=25, device=None, sample_rate=16000): |
| device = device or self.device |
| audio_duration = len(speech_array) / sample_rate |
| video_length = max(int(round(audio_duration * fps)), 1) |
| encoder_length = max(int(math.ceil(audio_duration * 50)), 1) |
| speech_array = self._loudness_norm(speech_array, sample_rate) |
| speech_array = self._add_noise_floor(speech_array) |
| speech_array = self._smooth_transients(speech_array, sample_rate) |
|
|
| layer_groups = [(0, 8), (8, 16), (16, 24), (24, 32), (32, 33)] |
| audio_chunks = [[] for _ in layer_groups] |
| mel_chunk = 750 * 640 |
| for start in range(0, len(speech_array), mel_chunk): |
| chunk = speech_array[start : start + mel_chunk] |
| features = self.audio_feature_extractor( |
| [chunk], |
| return_tensors="pt", |
| return_attention_mask=True, |
| sampling_rate=sample_rate, |
| ) |
| input_features = features.input_features.to(device=device, dtype=self.dtype) |
| attention_mask = getattr(features, "attention_mask", None) |
| if attention_mask is not None: |
| attention_mask = attention_mask.to(device=device) |
| input_features = self.audio_encoder._mask_input_features(input_features, attention_mask=attention_mask) |
| encoder_dtype = getattr(self.audio_encoder.encoder, "dtype", self.dtype) |
| outputs = self.audio_encoder.encoder( |
| input_features.to(encoder_dtype), |
| head_mask=None, |
| output_attentions=False, |
| output_hidden_states=True, |
| return_dict=True, |
| ) |
| chunk_states = outputs.hidden_states |
| for group_idx, (layer_start, layer_end) in enumerate(layer_groups): |
| layer_state = torch.stack(chunk_states[layer_start:layer_end], dim=0).mean(dim=0).squeeze(0) |
| audio_chunks[group_idx].append(layer_state.detach().to("cpu")) |
| del outputs, chunk_states, input_features, attention_mask, features |
|
|
| audio_layers = [] |
| for chunks in audio_chunks: |
| layer_state = torch.cat(chunks, dim=0)[:encoder_length] |
| layer_state = self._interpolate_audio_state(layer_state, video_length).to(dtype=self.dtype) |
| audio_layers.append(layer_state) |
| audio_emb = torch.stack(audio_layers, dim=1).contiguous() |
| return torch.nan_to_num(audio_emb, nan=0.0, posinf=0.0, neginf=0.0) |
|
|
| @torch.no_grad() |
| def _get_audio_embedding(self, speech_array, fps=32, device="cpu", sample_rate=16000): |
| if self.is_avatar_v1_5: |
| return self._get_audio_embedding_whisper(speech_array, fps=fps, device=device, sample_rate=sample_rate) |
| return self._get_audio_embedding_wav2vec(speech_array, fps=fps, device=device, sample_rate=sample_rate) |
|
|
| def _build_audio_windows(self, audio_path, frame_num, fps, window_start_frame_no, audio_stride): |
| speech_array, sr = librosa.load(audio_path, sr=16000) |
| target_len = int((window_start_frame_no + frame_num) / fps * sr) |
| if len(speech_array) < target_len: |
| pad = target_len - len(speech_array) |
| speech_array = np.pad(speech_array, (0, pad), mode="constant") |
|
|
| audio_device = self.device if self.is_avatar_v1_5 else "cpu" |
| full_audio_emb = self._get_audio_embedding(speech_array, fps=fps * audio_stride, device=audio_device, sample_rate=sr) |
| if torch.isnan(full_audio_emb).any(): |
| raise ValueError("Audio embedding contains NaNs.") |
|
|
| audio_start_idx = window_start_frame_no * audio_stride |
| audio_end_idx = audio_start_idx + audio_stride * frame_num |
| window = self.transformer.audio_window if hasattr(self.transformer, "audio_window") else 5 |
| offsets = torch.arange(window, device=full_audio_emb.device) - window // 2 |
| centers = torch.arange(audio_start_idx, audio_end_idx, audio_stride, device=full_audio_emb.device).unsqueeze(1) + offsets.unsqueeze(0) |
| centers = torch.clamp(centers, min=0, max=full_audio_emb.shape[0] - 1) |
| audio_emb = full_audio_emb[centers][None, ...].to("cpu") |
| del full_audio_emb |
| return audio_emb |
|
|
| def _build_ref_target_masks(self, height, width, speakers_bboxes=None): |
| if not speakers_bboxes: |
| speakers_bboxes = {"person1": [5, 10, 45, 90], "person2": [55, 10, 95, 90]} |
| human_masks = [] |
| background_mask = torch.zeros([height, width]) |
| for _, person_bbox in speakers_bboxes.items(): |
| y_min, x_min, y_max, x_max = person_bbox |
| x_min, y_min, x_max, y_max = max(x_min, 5), max(y_min, 5), min(x_max, 95), min(y_max, 95) |
| x_min, y_min, x_max, y_max = ( |
| int(height * x_min / 100), |
| int(width * y_min / 100), |
| int(height * x_max / 100), |
| int(width * y_max / 100), |
| ) |
| human_mask = torch.zeros([height, width]) |
| human_mask[int(x_min) : int(x_max), int(y_min) : int(y_max)] = 1 |
| background_mask += human_mask |
| human_masks.append(human_mask) |
| background_mask = torch.where(background_mask > 0, torch.tensor(0), torch.tensor(1)) |
| human_masks.append(background_mask) |
| return torch.stack(human_masks, dim=0) |
|
|
| def get_timesteps_sigmas(self, sampling_steps, use_distill=False): |
| if use_distill: |
| distill_indices = torch.arange(1, self.num_distill_sample_steps + 1, dtype=torch.float32) |
| distill_indices = (distill_indices * (self.num_timesteps // self.num_distill_sample_steps)).round().long() |
| if self.is_avatar_v1_5: |
| distill_indices = self.num_timesteps - distill_indices |
| sigmas = torch.flip(torch.linspace(0, 1, self.num_timesteps), [0]) |
| sigmas = torch.flip(sigmas[distill_indices], [0]).float() |
| if sampling_steps != self.num_distill_sample_steps: |
| inference_indices = np.linspace(0, self.num_distill_sample_steps, num=sampling_steps, endpoint=False) |
| sigmas = sigmas[np.floor(inference_indices).astype(np.int64)] |
| else: |
| inference_indices = np.linspace(0, self.num_distill_sample_steps, num=sampling_steps, endpoint=False) |
| inference_indices = np.floor(inference_indices).astype(np.int64) |
| sigmas = torch.flip(distill_indices, [0])[inference_indices].float() / self.num_timesteps |
| else: |
| sigmas = torch.linspace(1, 0.001, sampling_steps, dtype=torch.float32) |
| return sigmas.to(dtype=torch.float32, device="cpu") |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| seed=None, |
| input_prompt="", |
| n_prompt="", |
| sampling_steps=50, |
| input_ref_images=None, |
| input_frames=None, |
| input_frames2=None, |
| input_masks=None, |
| input_masks2=None, |
| input_video=None, |
| image_start=None, |
| image_end=None, |
| input_ref_masks=None, |
| input_faces=None, |
| input_custom=None, |
| frame_num=93, |
| batch_size=1, |
| height=480, |
| width=832, |
| fit_into_canvas=None, |
| alt_prompt=None, |
| guide_scale=4.0, |
| guide2_scale=None, |
| guide3_scale=None, |
| shift=None, |
| audio_cfg_scale=None, |
| joint_pass=False, |
| VAE_tile_size=None, |
| prefix_frames_count=0, |
| conditioning_latents_size=0, |
| callback=None, |
| embedded_guidance_scale=None, |
| enable_RIFLEx=None, |
| cfg_star_switch=False, |
| cfg_zero_step=-1, |
| apg_switch=None, |
| perturbation_switch=None, |
| perturbation_layers=None, |
| perturbation_start=None, |
| perturbation_end=None, |
| switch_threshold=None, |
| switch2_threshold=None, |
| guide_phases=None, |
| model_switch_phase=None, |
| alt_guide_scale=None, |
| input_waveform=None, |
| input_waveform_sample_rate=None, |
| audio_guide=None, |
| audio_guide2=None, |
| audio_prompt_type=None, |
| audio_proj=None, |
| audio_scale=None, |
| audio_context_lens=None, |
| context_scale=None, |
| control_scale_alt=None, |
| alt_scale=None, |
| motion_amplitude=None, |
| model_mode=None, |
| causal_block_size=None, |
| causal_attention=None, |
| fps=None, |
| window_start_frame_no=0, |
| sample_solver=None, |
| reference_image_enabled=None, |
| ref_img_index=None, |
| mask_frame_range=None, |
| overlapped_latents=None, |
| return_latent_slice=None, |
| speakers_bboxes=None, |
| window_no=None, |
| overlap_noise=None, |
| overlap_size=None, |
| color_correction_strength=None, |
| input_video_is_hdr=False, |
| lora_dir=None, |
| keep_frames_parsed=None, |
| model_filename=None, |
| model_type=None, |
| loras_slists=None, |
| NAG_scale=None, |
| NAG_tau=None, |
| NAG_alpha=None, |
| image_mode=None, |
| video_prompt_type=None, |
| offloadobj=None, |
| set_header_text=None, |
| pre_video_frame=None, |
| prefix_video=None, |
| original_input_ref_images=None, |
| image_refs_relative_size=None, |
| outpainting_dims=None, |
| face_arc_embeds=None, |
| custom_settings=None, |
| **kwargs, |
| ): |
| if self._interrupt: |
| return None |
|
|
| if seed is None or seed == -1: |
| seed = torch.seed() % (2**32 - 1) |
| generator = torch.Generator(device=self.device) |
| generator.manual_seed(seed) |
|
|
| if fps is None or fps == 0: |
| fps = self.model_def.get("fps", 15 if not self.is_avatar else 16) |
|
|
| if frame_num % self.vae_scale_factor_temporal != 1: |
| frame_num = frame_num // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 |
| frame_num = max(frame_num, 1) |
|
|
| sample_solver = self.model_def.get("sample_solver", "auto") if sample_solver is None else sample_solver |
| if sample_solver in (None, ""): |
| sample_solver = "default" |
|
|
| prompt_embeds, prompt_mask, neg_embeds, neg_mask = self._encode_prompt( |
| input_prompt, |
| n_prompt if n_prompt is not None else "", |
| device=self.device, |
| dtype=self.dtype, |
| ) |
|
|
| any_guidance = guide_scale is not None and guide_scale > 1 |
| if self.is_avatar: |
| if audio_cfg_scale is None: |
| audio_cfg_scale = 1.0 |
| any_guidance = any_guidance or audio_cfg_scale > 1 |
|
|
| if reference_image_enabled is None: |
| reference_image_enabled = self.model_def.get("reference_image_enabled", True) |
| reference_image_enabled = self.is_avatar and bool(reference_image_enabled) |
| reference_features_enabled = reference_image_enabled |
|
|
| ref_img_index = self.model_def.get("ref_img_index", 10) if ref_img_index is None else ref_img_index |
| mask_frame_range = self.model_def.get("mask_frame_range", 3) if mask_frame_range is None else mask_frame_range |
| if not reference_features_enabled: |
| ref_img_index = None |
| mask_frame_range = None |
|
|
| ref_image = None |
| if reference_image_enabled: |
| if input_ref_images is not None: |
| ref_list = input_ref_images if isinstance(input_ref_images, list) else [input_ref_images] |
| if len(ref_list) > 0: |
| ref_image = ref_list[0] |
| if window_no == 1: |
| if ref_image is not None: |
| self._reference_image = ( |
| ref_image.detach().to("cpu") if torch.is_tensor(ref_image) else ref_image |
| ) |
| else: |
| self._reference_image = None |
| if ref_image is None and self._reference_image is not None: |
| ref_image = self._reference_image |
|
|
| cond_video = None |
| num_cond_frames = 0 |
| if input_video is not None: |
| cond_video = input_video |
| num_cond_frames = max(int(prefix_frames_count or 0), 0) |
|
|
| self._apply_vae_tiling(VAE_tile_size) |
|
|
| if cond_video is not None: |
| cond_video = cond_video.to(device=self.device, dtype=self.VAE_dtype) |
| if cond_video.dim() == 4: |
| cond_video = cond_video.unsqueeze(0) |
| cond_video_frames = cond_video.shape[2] |
| if num_cond_frames <= 0: |
| cond_video = None |
| num_cond_frames = 0 |
| else: |
| num_cond_frames = min(num_cond_frames, cond_video_frames) |
|
|
| if sample_solver not in ("auto", "default", "enhance_hf", "distill"): |
| raise ValueError(f"Unsupported scheduler '{sample_solver}' for LongCat.") |
| if self.model_def.get("distill_only", False) and sample_solver != "distill": |
| raise ValueError("LongCat Avatar 1.5 currently supports the distilled scheduler only.") |
| use_distill = sample_solver == "distill" |
| enhance_hf = sample_solver == "enhance_hf" |
| if sample_solver == "auto": |
| enhance_hf = cond_video is not None and num_cond_frames > 1 |
| if use_distill and enhance_hf: |
| raise ValueError("distill and enhance_hf schedules cannot both be enabled.") |
|
|
| image_cond = None |
| ref_latent = None |
| num_ref_latents = 0 |
| if reference_image_enabled and ref_image is not None: |
| if not torch.is_tensor(ref_image): |
| ref_image = torch.from_numpy(np.array(ref_image)).float().div_(127.5).sub_(1.).movedim(-1, 0) |
| ref_image = ref_image.to(device=self.device, dtype=self.VAE_dtype) |
| if ref_image.dim() == 3: |
| ref_image = ref_image.unsqueeze(0) |
| if ref_image.dim() == 5 and ref_image.shape[2] == 1: |
| ref_image = ref_image.squeeze(2) |
| if ref_image.dim() != 4: |
| raise ValueError("reference image must be CHW or BCHW for LongCat.") |
| if ref_image.shape[0] == 1 and batch_size > 1: |
| ref_image = ref_image.repeat(batch_size, 1, 1, 1) |
| elif ref_image.shape[0] != batch_size: |
| raise ValueError("reference image batch size does not match prompts.") |
| if cond_video is None: |
| image_cond = ref_image |
| else: |
| ref_image_5d = ref_image.unsqueeze(2) |
| ref_latent = retrieve_latents(self.vae.encode(ref_image_5d), generator, sample_mode="argmax") |
| ref_latent = self.normalize_latents(ref_latent).to(torch.float32) |
| num_ref_latents = 1 |
|
|
| if torch.is_tensor(overlapped_latents): |
| if overlapped_latents.dim() == 4: |
| overlapped_latents = overlapped_latents.unsqueeze(0) |
| if overlapped_latents.dim() != 5: |
| overlapped_latents = None |
| else: |
| overlapped_latents = None |
|
|
| cond_image_frames = 1 if image_cond is not None else num_cond_frames |
| expected_num_cond_latents = ( |
| 1 + (cond_image_frames - 1) // self.vae_scale_factor_temporal if cond_image_frames > 0 else 0 |
| ) |
| use_overlap_latents = ( |
| overlapped_latents is not None and expected_num_cond_latents > 0 and image_cond is None |
| ) |
| if use_overlap_latents: |
| lat_h = int(height) // self.vae_scale_factor_spatial |
| lat_w = int(width) // self.vae_scale_factor_spatial |
| if ( |
| overlapped_latents.shape[1] != self.transformer.config.in_channels |
| or overlapped_latents.shape[3] != lat_h |
| or overlapped_latents.shape[4] != lat_w |
| ): |
| use_overlap_latents = False |
|
|
| if use_overlap_latents: |
| num_latent_frames = (frame_num - 1) // self.vae_scale_factor_temporal + 1 |
| shape = ( |
| batch_size, |
| self.transformer.config.in_channels, |
| num_latent_frames, |
| lat_h, |
| lat_w, |
| ) |
| latents = torch.randn(shape, generator=generator, device=self.device, dtype=torch.float32) |
|
|
| overlap_latents = overlapped_latents.to(device=self.device, dtype=torch.float32) |
| if overlap_latents.shape[0] == 1 and batch_size > 1: |
| overlap_latents = overlap_latents.repeat(batch_size, 1, 1, 1, 1) |
| if overlap_latents.shape[2] > expected_num_cond_latents: |
| overlap_latents = overlap_latents[:, :, -expected_num_cond_latents:] |
|
|
| cond_latents = None |
| if cond_video is not None and overlap_latents.shape[2] < expected_num_cond_latents: |
| cond_latents_list = [] |
| for i in range(batch_size): |
| encoded_input = cond_video[i][:, -cond_image_frames:].unsqueeze(0) |
| latent = retrieve_latents( |
| self.vae.encode(encoded_input), |
| generator, |
| sample_mode="argmax", |
| ) |
| cond_latents_list.append(latent) |
| cond_latents = torch.cat(cond_latents_list, dim=0).to(torch.float32) |
| cond_latents = self.normalize_latents(cond_latents) |
| overlap_len = min(overlap_latents.shape[2], cond_latents.shape[2]) |
| if overlap_len > 0: |
| cond_latents[:, :, -overlap_len:] = overlap_latents[:, :, -overlap_len:] |
| else: |
| cond_latents = overlap_latents |
|
|
| num_cond_latents = min(cond_latents.shape[2], num_latent_frames) if cond_latents is not None else 0 |
| if num_cond_latents > 0: |
| latents[:, :, :num_cond_latents] = cond_latents[:, :, -num_cond_latents:] |
| else: |
| latents, num_cond_latents = self._prepare_latents( |
| batch_size=batch_size, |
| num_channels_latents=self.transformer.config.in_channels, |
| height=height, |
| width=width, |
| num_frames=frame_num, |
| dtype=torch.float32, |
| device=self.device, |
| generator=generator, |
| latents=None, |
| image=image_cond, |
| video=None if image_cond is not None else cond_video, |
| num_cond_frames=cond_image_frames, |
| ) |
| if reference_image_enabled and ref_latent is None and self.is_avatar and num_cond_latents > 1: |
| ref_latent = latents[:, :, :1].clone() |
| num_ref_latents = 1 |
| if ref_latent is not None: |
| num_cond_latents += num_ref_latents |
| latents = torch.cat([ref_latent, latents], dim=2) |
| self._validate_sparse_attention(latents) |
|
|
| sigmas = self.get_timesteps_sigmas(sampling_steps, use_distill=use_distill) |
| self.scheduler.set_timesteps(sampling_steps, sigmas=sigmas, device=self.device) |
| timesteps = self.scheduler.timesteps |
| if enhance_hf: |
| num_tail_uniform_steps = max(3, min(15, int(len(timesteps) * 0.2))) |
| tail_uniform_start = float(timesteps.max()) * 0.5 |
| tail_uniform_end = 0 |
| timesteps_uniform_tail = list( |
| np.linspace( |
| tail_uniform_start, |
| tail_uniform_end, |
| num_tail_uniform_steps, |
| dtype=np.float32, |
| endpoint=(tail_uniform_end != 0), |
| ) |
| ) |
| timesteps_uniform_tail = [ |
| torch.tensor(t, device=self.device, dtype=torch.float32).unsqueeze(0) |
| for t in timesteps_uniform_tail |
| ] |
| filtered_timesteps = [ |
| timestep.unsqueeze(0).to(self.device) for timestep in timesteps if timestep > tail_uniform_start |
| ] |
| timesteps = torch.cat(filtered_timesteps + timesteps_uniform_tail) |
| self.scheduler.timesteps = timesteps |
| self.scheduler.sigmas = torch.cat( |
| [timesteps / self.num_timesteps, torch.zeros(1, device=timesteps.device)] |
| ) |
|
|
| audio_emb = None |
| ref_target_masks = None |
| if self.is_avatar: |
| if audio_guide is None: |
| raise ValueError("Audio guide is required for LongCat Avatar.") |
| audio_stride = int(self.model_def.get("audio_stride", 1 if self.is_avatar_v1_5 else 2)) |
| audio_emb = self._build_audio_windows( |
| audio_guide, frame_num, fps, window_start_frame_no, audio_stride |
| ) |
| if audio_guide2 is not None or self.model_def.get("multi_speakers_only", False): |
| if audio_guide2 is None: |
| raise ValueError("Second audio guide is required for LongCat Avatar Multi.") |
| audio_emb2 = self._build_audio_windows( |
| audio_guide2, frame_num, fps, window_start_frame_no, audio_stride |
| ) |
| audio_emb = torch.cat([audio_emb, audio_emb2], dim=0) |
| ref_target_masks = self._build_ref_target_masks(height, width, speakers_bboxes) |
| if ref_target_masks is not None: |
| ref_target_masks = ref_target_masks.to(self.device) |
| if self.is_avatar_v1_5 and offloadobj is not None: |
| offloadobj.unload_all() |
| audio_emb = audio_emb.to(self.device, dtype=self.dtype) |
|
|
| latents = latents.to(self.device, dtype=self.dtype) |
| prompt_embeds = prompt_embeds.to(self.device) |
| prompt_mask = prompt_mask.to(self.device) |
| if neg_embeds is None: |
| neg_embeds = prompt_embeds |
| neg_mask = prompt_mask |
| else: |
| neg_embeds = neg_embeds.to(self.device) |
| neg_mask = neg_mask.to(self.device) |
|
|
| ref_kwargs = {} |
| if self.is_avatar and num_ref_latents > 0: |
| ref_kwargs = { |
| "num_ref_latents": num_ref_latents, |
| "ref_img_index": ref_img_index, |
| "mask_frame_range": mask_frame_range, |
| } |
|
|
| callback(-1, None, True, override_num_inference_steps = len(timesteps)) |
|
|
| with tqdm(total=len(timesteps), desc="Denoising") as progress_bar: |
| for i, t in enumerate(timesteps): |
| if self._interrupt: |
| return None |
| def _aborted(outputs): |
| if outputs is None: |
| return True |
| if isinstance(outputs, (list, tuple)): |
| return any(item is None for item in outputs) |
| return outputs is None |
|
|
| timestep = t.expand(latents.shape[0]).to(self.dtype) |
| if num_cond_latents > 0: |
| timestep = timestep[:, None].expand(-1, latents.shape[2]).clone() |
| timestep[:, :num_cond_latents] = 0 |
|
|
| if self.is_avatar and audio_emb is not None and any_guidance: |
| audio_cond = audio_emb.to(self.device, dtype=self.dtype) |
| audio_uncond = torch.zeros_like(audio_cond) |
| x_list = [latents, latents, latents] |
| ctx_list = [prompt_embeds, neg_embeds, neg_embeds] |
| mask_list = [prompt_mask, neg_mask, neg_mask] |
| audio_list = [audio_cond, audio_cond, audio_uncond] |
| ref_list = [ref_target_masks, ref_target_masks, ref_target_masks] |
| if joint_pass: |
| outputs = self.transformer( |
| hidden_states=x_list, |
| timestep=[timestep] * len(x_list), |
| encoder_hidden_states=ctx_list, |
| encoder_attention_mask=mask_list, |
| num_cond_latents=[num_cond_latents] * len(x_list), |
| audio_embs=audio_list, |
| ref_target_masks=ref_list, |
| **ref_kwargs, |
| ) |
| if _aborted(outputs): |
| return None |
| else: |
| outputs = [] |
| for x_i, ctx_i, mask_i, audio_i, ref_i in zip( |
| x_list, ctx_list, mask_list, audio_list, ref_list |
| ): |
| output = self.transformer( |
| hidden_states=x_i, |
| timestep=timestep, |
| encoder_hidden_states=ctx_i, |
| encoder_attention_mask=mask_i, |
| num_cond_latents=num_cond_latents, |
| audio_embs=audio_i, |
| ref_target_masks=ref_i, |
| **ref_kwargs, |
| ) |
| if _aborted(output): |
| return None |
| outputs.append(output) |
| noise_pred_cond, noise_pred_uncond_text, noise_pred_uncond = outputs |
| noise_pred = ( |
| noise_pred_uncond |
| + guide_scale * (noise_pred_cond - noise_pred_uncond_text) |
| + audio_cfg_scale * (noise_pred_uncond_text - noise_pred_uncond) |
| ) |
| elif any_guidance: |
| x_list = [latents, latents] |
| ctx_list = [prompt_embeds, neg_embeds] |
| mask_list = [prompt_mask, neg_mask] |
| if joint_pass: |
| outputs = self.transformer( |
| hidden_states=x_list, |
| timestep=[timestep] * len(x_list), |
| encoder_hidden_states=ctx_list, |
| encoder_attention_mask=mask_list, |
| num_cond_latents=[num_cond_latents] * len(x_list), |
| **ref_kwargs, |
| ) |
| if _aborted(outputs): |
| return None |
| else: |
| outputs = [] |
| for x_i, ctx_i, mask_i in zip(x_list, ctx_list, mask_list): |
| output = self.transformer( |
| hidden_states=x_i, |
| timestep=timestep, |
| encoder_hidden_states=ctx_i, |
| encoder_attention_mask=mask_i, |
| num_cond_latents=num_cond_latents, |
| **ref_kwargs, |
| ) |
| if _aborted(output): |
| return None |
| outputs.append(output) |
| noise_pred_cond, noise_pred_uncond = outputs |
| if cfg_star_switch: |
| positive_flat = noise_pred_cond.view(latents.shape[0], -1) |
| negative_flat = noise_pred_uncond.view(latents.shape[0], -1) |
| st_star = optimized_scale(positive_flat, negative_flat).view(latents.shape[0], 1, 1, 1) |
| if cfg_zero_step >= 0 and i <= cfg_zero_step: |
| noise_pred = noise_pred_cond * 0.0 |
| else: |
| noise_pred_uncond = noise_pred_uncond * st_star |
| noise_pred = noise_pred_uncond + guide_scale * (noise_pred_cond - noise_pred_uncond) |
| else: |
| noise_pred = noise_pred_uncond + guide_scale * (noise_pred_cond - noise_pred_uncond) |
| else: |
| noise_pred = self.transformer( |
| hidden_states=latents, |
| timestep=timestep, |
| encoder_hidden_states=prompt_embeds, |
| encoder_attention_mask=prompt_mask, |
| num_cond_latents=num_cond_latents, |
| audio_embs=audio_emb if self.is_avatar else None, |
| ref_target_masks=ref_target_masks if self.is_avatar else None, |
| **ref_kwargs, |
| ) |
| if _aborted(noise_pred): |
| return None |
|
|
| noise_pred = -noise_pred |
|
|
| if num_cond_latents > 0: |
| latents[:, :, num_cond_latents:] = self.scheduler.step( |
| noise_pred[:, :, num_cond_latents:], |
| t, |
| latents[:, :, num_cond_latents:], |
| return_dict=False, |
| )[0] |
| else: |
| latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] |
|
|
| if callback is not None: |
| callback(i, latents.squeeze(0)) |
| progress_bar.update() |
|
|
| if num_ref_latents > 0: |
| latents = latents[:, :, num_ref_latents:] |
| num_cond_latents -= num_ref_latents |
|
|
| latent_slice = None |
| if return_latent_slice is not None: |
| latent_slice = latents[:, :, return_latent_slice].detach().to("cpu") |
|
|
| latents = latents.to(self.vae.dtype) |
| latents = self.denormalize_latents(latents) |
| video = self.vae.decode(latents, return_dict=False)[0].clamp(-1, 1) |
| if video.dim() == 5: |
| video = video[0] |
| self._clear_runtime_caches() |
|
|
| if latent_slice is not None: |
| return {"x": video, "latent_slice": latent_slice} |
| return video |
|
|
