273 lines
10 KiB
Python
273 lines
10 KiB
Python
#!/app/miniconda3/envs/comfyui/bin/python
|
|
import os
|
|
import sys
|
|
|
|
# Paranoid import checks
|
|
print("--- Importing Modules ---")
|
|
try:
|
|
import torch
|
|
print(f"Torch imported: {torch.__version__}")
|
|
except ImportError as e:
|
|
print(f"Failed to import torch: {e}")
|
|
|
|
try:
|
|
import torchaudio
|
|
print(f"Torchaudio imported: {torchaudio.__version__}")
|
|
except ImportError as e:
|
|
print(f"FATAL: Failed to import torchaudio: {e}")
|
|
sys.exit(1)
|
|
|
|
import numpy as np
|
|
import math
|
|
import gc
|
|
import logging
|
|
from PIL import Image
|
|
from omegaconf import OmegaConf
|
|
|
|
# Mock ComfyUI environment
|
|
class MockFolderPaths:
|
|
models_dir = "/app/ComfyUI/models"
|
|
|
|
folder_paths = MockFolderPaths()
|
|
|
|
# Configure Paths
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
root_dir = os.path.dirname(current_dir) # ComfyUI_AIIA root
|
|
echomimic_v3_root = os.path.join(root_dir, "libs", "EchoMimicV3")
|
|
|
|
sys.path.insert(0, root_dir)
|
|
sys.path.insert(0, echomimic_v3_root)
|
|
|
|
# Imports from EchoMimicV3
|
|
print("--- Importing EchoMimicV3 Modules ---")
|
|
from transformers import AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
from echomimic_v3_src.wan_vae import AutoencoderKLWan
|
|
from echomimic_v3_src.wan_image_encoder import CLIPModel
|
|
from echomimic_v3_src.wan_text_encoder import WanT5EncoderModel
|
|
from echomimic_v3_src.wan_transformer3d_audio import WanTransformerAudioMask3DModel
|
|
from echomimic_v3_src.pipeline_wan_fun_inpaint_audio import WanFunInpaintAudioPipeline
|
|
from echomimic_v3_src.utils import get_image_to_video_latent3, filter_kwargs
|
|
from echomimic_v3_src.face_detect import get_mask_coord
|
|
import torchvision.transforms.functional as TF
|
|
from echomimic_v3_src.fm_solvers import FlowDPMSolverMultistepScheduler
|
|
|
|
# Config
|
|
ECHOMIMIC_MODELS_DIR = "EchoMimicV3"
|
|
MODEL_SUBFOLDER = "EchoMimicV3"
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
DTYPE = torch.bfloat16
|
|
SAMPLE_IMAGE_PATH = "/app/ComfyUI/input/xuerOneCyanTenColor_fluxV10--20241112-194257-00001.png"
|
|
SAMPLE_AUDIO_PATH = "/app/ComfyUI/custom_nodes/ComfyUI_AIIA/assets/seed_male.wav"
|
|
|
|
def get_ip_mask(coords):
|
|
y1, y2, x1, x2, h, w = coords
|
|
Y, X = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij')
|
|
mask = (Y.unsqueeze(-1) >= y1) & (Y.unsqueeze(-1) < y2) & (X.unsqueeze(-1) >= x1) & (X.unsqueeze(-1) < x2)
|
|
mask = mask.reshape(-1)
|
|
return mask.float()
|
|
|
|
def main():
|
|
print(f"=== Starting EchoMimicV3 SERVER FINAL Test ===")
|
|
print(f"Torchaudio available check: {torchaudio}")
|
|
print(f"Device: {DEVICE}, DType: {DTYPE}")
|
|
|
|
# 1. Load Config
|
|
config_path = os.path.join(echomimic_v3_root, "config", "config.yaml")
|
|
cfg = OmegaConf.load(config_path)
|
|
|
|
# 2. Paths
|
|
models_base = os.path.join(folder_paths.models_dir, ECHOMIMIC_MODELS_DIR)
|
|
|
|
# Specific paths based on 'ls -R' output from server
|
|
echomimic_root = os.path.join(models_base, "EchoMimicV3")
|
|
wan_root = os.path.join(models_base, "Wan2.1-Fun-V1.1-1.3B-InP")
|
|
|
|
print(f"EchoMimic Root: {echomimic_root}")
|
|
print(f"Wan Root: {wan_root}")
|
|
|
|
# 3. Load Models
|
|
print("Loading Transformer...")
|
|
# Transformer is explicitly in EchoMimicV3/transformer
|
|
transformer_path = os.path.join(echomimic_root, "transformer")
|
|
if not os.path.exists(transformer_path):
|
|
print(f"Warning: Transformer path not found at {transformer_path}, checking Wan root")
|
|
transformer_path = wan_root
|
|
|
|
print(f"Loading transformer from: {transformer_path}")
|
|
transformer = WanTransformerAudioMask3DModel.from_pretrained(
|
|
transformer_path,
|
|
transformer_additional_kwargs=OmegaConf.to_container(cfg['transformer_additional_kwargs']),
|
|
torch_dtype=torch.float32, # Load as float32 then move/cast
|
|
low_cpu_mem_usage=True
|
|
).to("cpu").to(DTYPE) # Keep on CPU first
|
|
|
|
print("Loading VAE...")
|
|
# VAE is in Wan2.1 folder
|
|
vae = AutoencoderKLWan.from_pretrained(
|
|
os.path.join(wan_root, "Wan2.1_VAE.pth"),
|
|
additional_kwargs=OmegaConf.to_container(cfg['vae_kwargs']),
|
|
).to(dtype=torch.float32, device="cpu") # VAE on CPU initially
|
|
|
|
print("Loading Tokenizer...")
|
|
tokenizer = AutoTokenizer.from_pretrained(os.path.join(wan_root, "google/umt5-xxl"))
|
|
|
|
print("Loading Text Encoder...")
|
|
# WanT5EncoderModel expects a .pth file, not a directory or subfolder arg
|
|
text_encoder_file = os.path.join(wan_root, "models_t5_umt5-xxl-enc-bf16.pth")
|
|
print(f"Loading Text Encoder from: {text_encoder_file}")
|
|
|
|
if not os.path.exists(text_encoder_file):
|
|
raise FileNotFoundError(f"Text Encoder file not found: {text_encoder_file}")
|
|
|
|
text_encoder = WanT5EncoderModel.from_pretrained(
|
|
text_encoder_file,
|
|
additional_kwargs=OmegaConf.to_container(cfg['text_encoder_kwargs']),
|
|
torch_dtype=DTYPE,
|
|
low_cpu_mem_usage=True
|
|
).to(dtype=DTYPE, device="cpu").eval()
|
|
|
|
print("Loading Image Encoder...")
|
|
# Similarly, CLIPModel likely expects the .pth file
|
|
image_encoder_file = os.path.join(wan_root, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
|
|
print(f"Loading Image Encoder from: {image_encoder_file}")
|
|
|
|
if not os.path.exists(image_encoder_file):
|
|
raise FileNotFoundError(f"Image Encoder file not found: {image_encoder_file}")
|
|
|
|
clip_image_encoder = CLIPModel.from_pretrained(
|
|
image_encoder_file
|
|
).to(dtype=DTYPE, device="cpu").eval()
|
|
|
|
print("Loading Scheduler...")
|
|
scheduler_kwargs = OmegaConf.to_container(cfg['scheduler_kwargs'])
|
|
scheduler = FlowDPMSolverMultistepScheduler(**filter_kwargs(FlowDPMSolverMultistepScheduler, scheduler_kwargs))
|
|
|
|
print("Loading Audio Encoder...")
|
|
wav2vec_path = os.path.join(folder_paths.models_dir, ECHOMIMIC_MODELS_DIR, "wav2vec2-base-960h")
|
|
wav2vec_processor = Wav2Vec2Processor.from_pretrained(wav2vec_path)
|
|
wav2vec_model = Wav2Vec2Model.from_pretrained(wav2vec_path).to(dtype=DTYPE, device="cpu").eval()
|
|
|
|
# 4. Pipeline
|
|
pipeline = WanFunInpaintAudioPipeline(
|
|
transformer=transformer,
|
|
vae=vae,
|
|
tokenizer=tokenizer,
|
|
text_encoder=text_encoder,
|
|
scheduler=scheduler,
|
|
clip_image_encoder=clip_image_encoder,
|
|
).to("cpu")
|
|
|
|
# 5. Prepare Inputs
|
|
print("Preparing Inputs...")
|
|
if not os.path.exists(SAMPLE_IMAGE_PATH):
|
|
print(f"Image not found at {SAMPLE_IMAGE_PATH}, creating dummy...")
|
|
ref_img_pil = Image.new("RGB", (768, 768), (100, 100, 200))
|
|
else:
|
|
ref_img_pil = Image.open(SAMPLE_IMAGE_PATH).convert("RGB")
|
|
ref_img_pil = ref_img_pil.resize((768, 768)) # Force resize
|
|
|
|
# Audio
|
|
if not os.path.exists(SAMPLE_AUDIO_PATH):
|
|
raise FileNotFoundError(f"Audio not found: {SAMPLE_AUDIO_PATH}")
|
|
|
|
print(f"Loading Audio from {SAMPLE_AUDIO_PATH}")
|
|
audio_wav, sr = torchaudio.load(SAMPLE_AUDIO_PATH)
|
|
if sr != 16000:
|
|
resampler = torchaudio.transforms.Resample(sr, 16000)
|
|
audio_wav = resampler(audio_wav)
|
|
|
|
# Process Audio
|
|
# Simplified audio feature extraction
|
|
audio_inputs = wav2vec_processor(audio_wav[0].numpy(), sampling_rate=16000, return_tensors="pt").input_values
|
|
with torch.no_grad():
|
|
audio_embeds = wav2vec_model(audio_inputs.to(dtype=DTYPE)).last_hidden_state
|
|
audio_embeds = audio_embeds.to(DEVICE, dtype=DTYPE)
|
|
|
|
# Face Mask (Simplified - full face)
|
|
# EchoMimicV3 usually detects face, here we just make a centered mask
|
|
print("Generating Mask...")
|
|
h, w = 768, 768
|
|
y1, y2, x1, x2 = h//4, h*3//4, w//4, w*3//4
|
|
ip_mask = get_ip_mask((y1, y2, x1, x2, h, w)).to(DEVICE, dtype=DTYPE).unsqueeze(0)
|
|
|
|
# 6. Run Generation (Short chunk)
|
|
print("Starting Generation...")
|
|
|
|
# Manual Memory Mgmt (Mimic Node)
|
|
|
|
# Encode Prompt
|
|
print("Encoding Prompt...")
|
|
pipeline.text_encoder.to(DEVICE)
|
|
prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt(
|
|
prompt="a talking head video",
|
|
negative_prompt="bad quality",
|
|
do_classifier_free_guidance=True,
|
|
num_videos_per_prompt=1,
|
|
max_sequence_length=512,
|
|
device=DEVICE
|
|
)
|
|
pipeline.text_encoder.to("cpu")
|
|
torch.cuda.empty_cache()
|
|
|
|
# Move Models
|
|
pipeline.transformer.to(DEVICE)
|
|
pipeline.vae.to(DEVICE)
|
|
pipeline.clip_image_encoder.to(DEVICE) # Ensure clip encoder is on device for context
|
|
|
|
# Prepare Latents
|
|
input_video, input_video_mask, clip_image = get_image_to_video_latent3(
|
|
ref_img_pil, None, video_length=25, sample_size=[768, 768]
|
|
)
|
|
|
|
# CLIP Context
|
|
print("Computing CLIP Context...")
|
|
# pipeline.clip_image_encoder.to(DEVICE) # Already moved
|
|
clip_image_t = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(DEVICE, dtype=DTYPE)
|
|
clip_context = pipeline.clip_image_encoder([clip_image_t[:, None, :, :]])
|
|
# pipeline.clip_image_encoder.to("cpu") # Keep on GPU for short test or move back
|
|
# For safety in this script, let's just keep everything on GPU if it fits
|
|
# But to test properly we should mimic node logic
|
|
pipeline.clip_image_encoder.to("cpu")
|
|
torch.cuda.empty_cache()
|
|
|
|
# Run
|
|
print("Running Pipeline Loop...")
|
|
partial_audio_embeds = audio_embeds[:, :50] # 25 frames * 2
|
|
|
|
with torch.no_grad():
|
|
sample = pipeline(
|
|
prompt=None,
|
|
num_frames=25,
|
|
negative_prompt=None,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
audio_embeds=partial_audio_embeds,
|
|
audio_scale=1.0,
|
|
ip_mask=ip_mask,
|
|
use_un_ip_mask=False,
|
|
height=768,
|
|
width=768,
|
|
generator=torch.Generator(device="cpu").manual_seed(42),
|
|
clip_context=clip_context,
|
|
neg_scale=1.5,
|
|
neg_steps=2,
|
|
use_dynamic_cfg=True,
|
|
use_dynamic_acfg=True,
|
|
guidance_scale=2.5,
|
|
audio_guidance_scale=1.0,
|
|
num_inference_steps=20, # Short run
|
|
video=input_video,
|
|
mask_video=input_video_mask,
|
|
clip_image=clip_image,
|
|
).videos
|
|
|
|
print("Generation/Inference successful!")
|
|
print(f"Output shape: {sample.shape}")
|
|
|
|
# Simple save to verify (Requires moviepy or opencv, let's just print success)
|
|
print("Test passed.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|