wan_runner.py

import subprocess
import shutil
import os
import sys
from pathlib import Path

BASE = Path(__file__).resolve().parent
CKPT_VACE = BASE / "Wan2.1-VACE-1.3B"
CKPT_FLF = BASE / "Wan2.1-FLF2V-14B-720P"

def run_cmd(cmd):
    """Run a command and handle errors"""
    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Command failed: {' '.join(cmd)}")
        print(f"Error: {e.stderr}")
        raise

def check_wan_installation():
    """Check if Wan2.1 and VACE are properly installed"""
    try:
        # Check if the repositories exist locally
        wan_path = BASE / "Wan2.1"
        vace_path = BASE / "VACE"
        
        if not wan_path.exists():
            print("Cloning Wan2.1 repository...")
            run_cmd(["git", "clone", "https://github.com/Wan-Video/Wan2.1.git", str(wan_path)])
        
        if not vace_path.exists():
            print("Cloning VACE repository...")
            run_cmd(["git", "clone", "https://github.com/ali-vilab/VACE.git", str(vace_path)])
        
        # Add to Python path
        sys.path.insert(0, str(wan_path))
        sys.path.insert(0, str(vace_path))
        
        return True
    except Exception as e:
        print(f"Failed to setup Wan2.1/VACE: {e}")
        return False

def generate_video(ref_img, first_frame, last_frame, output_path):
    """Generate video from reference image and first/last frames"""
    if not check_wan_installation():
        raise RuntimeError("Wan2.1/VACE setup failed")
    
    # Check if model directory exists
    if not os.path.exists(str(CKPT_FLF)):
        raise FileNotFoundError(f"モデルファイルが見つかりません: {CKPT_FLF}\n管理者に連絡してください。")
    
    # Create necessary directories
    os.makedirs("processed", exist_ok=True)
    os.makedirs("results", exist_ok=True)
    
    # Preprocess frames
    vace_script = BASE / "VACE" / "vace_preprocess.py"
    if vace_script.exists():
        run_cmd([
            sys.executable, str(vace_script),
            "--task", "extension",
            "--mode", "firstlastframe",
            "--image", f"{first_frame},{last_frame}",
            "--expand_num", "240"
        ])
    else:
        # Fallback: create a simple template video using imageio
        print("VACE preprocessing not available, using fallback method")
        import imageio
        import numpy as np
        from PIL import Image
        
        # Load first and last frames
        first_img = Image.open(first_frame).convert('RGB')
        last_img = Image.open(last_frame).convert('RGB')
        
        # Resize to target size
        target_size = (1280, 720)
        first_img = first_img.resize(target_size, Image.Resampling.LANCZOS)
        last_img = last_img.resize(target_size, Image.Resampling.LANCZOS)
        
        # Create video with interpolated gray frames
        frames = []
        frames.append(np.array(first_img))
        
        # Add gray frames in between
        gray_frame = np.ones((720, 1280, 3), dtype=np.uint8) * 128
        for _ in range(238):  # 240 total frames - 2 (first and last)
            frames.append(gray_frame)
        
        frames.append(np.array(last_img))
        
        # Save as video
        os.makedirs("processed", exist_ok=True)
        imageio.mimwrite("processed/src_video.mp4", frames, fps=24, codec='libx264')
    
    # Generate video
    generate_script = BASE / "Wan2.1" / "generate.py"
    if generate_script.exists():
        run_cmd([
            sys.executable, str(generate_script),
            "--task", "flf2v-14B",
            "--size", "1280*720",
            "--ckpt_dir", str(CKPT_FLF),
            "--src_ref_images", ref_img,
            "--src_video", "processed/src_video.mp4",
            "--first_frame", first_frame,
            "--last_frame", last_frame,
            "--offload_model", "True",
            "--t5_cpu"
        ])
        
        # Move result
        if os.path.exists("results/final.mp4"):
            shutil.move("results/final.mp4", output_path)
    else:
        raise FileNotFoundError("Wan2.1 generate.py not found")

def generate_image(ref_img, prompt, output_path):
    """Generate image from reference image and prompt"""
    if not check_wan_installation():
        raise RuntimeError("Wan2.1/VACE setup failed")
    
    # Check if model directory exists
    if not os.path.exists(str(CKPT_VACE)):
        raise FileNotFoundError(f"モデルファイルが見つかりません: {CKPT_VACE}\n管理者に連絡してください。")
    
    # Check for critical T5 model file
    t5_model_path = CKPT_VACE / "models_t5_umt5-xxl-enc-bf16.pth"
    if not t5_model_path.exists():
        raise FileNotFoundError(f"T5モデルファイルが見つかりません: {t5_model_path}\n管理者に連絡してください。")
    
    # Create necessary directories
    os.makedirs("results", exist_ok=True)
    
    # Generate image
    generate_script = BASE / "Wan2.1" / "generate.py"
    if generate_script.exists():
        run_cmd([
            sys.executable, str(generate_script),
            "--task", "vace-1.3B",
            "--size", "832*480",
            "--ckpt_dir", str(CKPT_VACE),
            "--src_ref_images", ref_img,
            "--frame_num", "1",
            "--prompt", prompt or " "
        ])
        
        # Move result
        if os.path.exists("results/final_frame000.png"):
            shutil.move("results/final_frame000.png", output_path)
    else:
        raise FileNotFoundError("Wan2.1 generate.py not found")