vace_integration.py

"""
VACEとWan2.1の統合ヘルパー
前処理と推論の実行を簡素化
"""
import os
import sys
import logging
import subprocess
from pathlib import Path
import numpy as np
import cv2
import imageio
from PIL import Image

logger = logging.getLogger(__name__)

class VACEProcessor:
    """VACE前処理と推論のラッパークラス"""
    
    def __init__(self, ckpt_dir):
        self.ckpt_dir = Path(ckpt_dir)
        self.check_installation()
    
    def check_installation(self):
        """VACEとWan2.1のインストール状態をチェック"""
        try:
            import wan
            import VACE
            logger.info("VACEとWan2.1が正常にインポートされました")
            return True
        except ImportError:
            logger.warning("VACEまたはWan2.1が見つかりません。フォールバックモードで動作します")
            return False
    
    def create_template_video(self, first_frame_path, last_frame_path, output_dir, num_frames=240):
        """
        最初と最後のフレームから補間用テンプレート動画を作成
        
        Args:
            first_frame_path: 最初のフレーム画像パス
            last_frame_path: 最後のフレーム画像パス
            output_dir: 出力ディレクトリ
            num_frames: 総フレーム数
        
        Returns:
            tuple: (動画パス, マスクパス)
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)
        
        # 画像を読み込み
        first_img = cv2.imread(str(first_frame_path))
        last_img = cv2.imread(str(last_frame_path))
        
        # サイズを統一（512x512）
        target_size = (512, 512)
        first_img = cv2.resize(first_img, target_size)
        last_img = cv2.resize(last_img, target_size)
        
        # テンプレート動画を作成
        video_frames = []
        mask_frames = []
        
        for i in range(num_frames):
            if i == 0:
                # 最初のフレーム
                video_frames.append(first_img)
                mask_frames.append(np.ones((512, 512), dtype=np.uint8) * 255)
            elif i == num_frames - 1:
                # 最後のフレーム
                video_frames.append(last_img)
                mask_frames.append(np.ones((512, 512), dtype=np.uint8) * 255)
            else:
                # 中間フレーム（グレー）
                gray_frame = np.ones((512, 512, 3), dtype=np.uint8) * 127
                video_frames.append(gray_frame)
                mask_frames.append(np.zeros((512, 512), dtype=np.uint8))
        
        # 動画として保存
        video_path = output_dir / "src_video.mp4"
        mask_path = output_dir / "src_mask.mp4"
        
        # ビデオライターを作成
        writer = imageio.get_writer(str(video_path), fps=24)
        mask_writer = imageio.get_writer(str(mask_path), fps=24)
        
        for frame, mask in zip(video_frames, mask_frames):
            # BGRからRGBに変換
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            writer.append_data(frame_rgb)
            
            # マスクを3チャンネルに拡張
            mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
            mask_writer.append_data(mask_3ch)
        
        writer.close()
        mask_writer.close()
        
        logger.info(f"テンプレート動画を作成しました: {video_path}")
        return str(video_path), str(mask_path)
    
    def run_inference(self, src_video, src_mask, ref_image, prompt, output_dir):
        """
        Wan2.1推論を実行（簡易版）
        
        実際の実装では、ここでWan2.1モデルを直接呼び出します
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)
        
        try:
            # 実際のWan2.1推論コマンドを実行
            cmd = [
                "python", "-m", "wan.inference",
                "--model_path", str(self.ckpt_dir),
                "--input_video", src_video,
                "--reference_image", ref_image,
                "--prompt", prompt,
                "--output_path", str(output_dir / "output.mp4"),
                "--num_inference_steps", "50",
                "--guidance_scale", "7.5"
            ]
            
            if src_mask:
                cmd.extend(["--mask_video", src_mask])
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                logger.error(f"推論エラー: {result.stderr}")
                # フォールバック: ダミー動画を作成
                return self.create_dummy_output(output_dir)
            
            return str(output_dir / "output.mp4")
            
        except Exception as e:
            logger.error(f"推論実行エラー: {e}")
            return self.create_dummy_output(output_dir)
    
    def create_dummy_output(self, output_dir):
        """開発/テスト用のダミー出力動画を作成"""
        output_path = output_dir / "output.mp4"
        
        # 簡単なテスト動画を作成（10秒、24fps）
        writer = imageio.get_writer(str(output_path), fps=24)
        
        for i in range(240):  # 10秒 × 24fps
            # グラデーション画像を作成
            frame = np.zeros((512, 512, 3), dtype=np.uint8)
            frame[:, :] = [
                int(255 * (i / 240)),  # R
                int(128),              # G  
                int(255 * (1 - i / 240))  # B
            ]
            writer.append_data(frame)
        
        writer.close()
        logger.info(f"ダミー動画を作成しました: {output_path}")
        return str(output_path)

def process_video_generation(ref_image_path, first_frame_path, last_frame_path, 
                           prompt, ckpt_dir, output_dir):
    """
    動画生成の全体プロセスを実行
    
    Args:
        ref_image_path: 参照画像パス
        first_frame_path: 最初のフレーム画像パス
        last_frame_path: 最後のフレーム画像パス
        prompt: テキストプロンプト
        ckpt_dir: モデルチェックポイントディレクトリ
        output_dir: 出力ディレクトリ
    
    Returns:
        str: 生成された動画のパス
    """
    processor = VACEProcessor(ckpt_dir)
    
    # 前処理: テンプレート動画作成
    processed_dir = Path(output_dir) / "processed"
    src_video, src_mask = processor.create_template_video(
        first_frame_path, 
        last_frame_path,
        processed_dir,
        num_frames=240
    )
    
    # 推論実行
    inference_dir = Path(output_dir) / "inference"
    output_video = processor.run_inference(
        src_video,
        src_mask,
        ref_image_path,
        prompt,
        inference_dir
    )
    
    return output_video

if __name__ == "__main__":
    # テスト実行
    import tempfile
    
    with tempfile.TemporaryDirectory() as temp_dir:
        # テスト画像を作成
        test_img = Image.new('RGB', (512, 512), color='red')
        test_img.save(f"{temp_dir}/test.png")
        
        # 処理実行
        output = process_video_generation(
            f"{temp_dir}/test.png",
            f"{temp_dir}/test.png", 
            f"{temp_dir}/test.png",
            "test prompt",
            "./Wan2.1-VACE-1.3B",
            temp_dir
        )
        
        print(f"Output video: {output}")