scripts/decode_openx.py

import os
import torch
import numpy as np
from PIL import Image
import imageio
import argparse
from diffsynth import WanVideoReCamMasterPipeline, ModelManager
from tqdm import tqdm
import json

class VideoDecoder:
    def __init__(self, vae_path, device="cuda"):
        """初始化视频解码器"""
        self.device = device
        
        # 初始化模型管理器
        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
        model_manager.load_models([vae_path])
        
        # 创建pipeline并只保留VAE
        self.pipe = WanVideoReCamMasterPipeline.from_model_manager(model_manager)
        self.pipe = self.pipe.to(device)
        
        # 🔧 关键修复：确保VAE及其所有组件都在正确设备上
        self.pipe.vae = self.pipe.vae.to(device)
        if hasattr(self.pipe.vae, 'model'):
            self.pipe.vae.model = self.pipe.vae.model.to(device)
        
        print(f"✅ VAE解码器初始化完成，设备: {device}")

    def decode_latents_to_video(self, latents, output_path, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)):
        """
        将latents解码为视频 - 修正版本，修复维度处理问题
        """
        print(f"🔧 开始解码latents...")
        print(f"输入latents形状: {latents.shape}")
        print(f"输入latents设备: {latents.device}")
        print(f"输入latents数据类型: {latents.dtype}")
        
        # 确保latents有batch维度
        if len(latents.shape) == 4:  # [C, T, H, W]
            latents = latents.unsqueeze(0)  # -> [1, C, T, H, W]
        
        # 🔧 关键修正：确保latents在正确的设备上且数据类型匹配
        model_dtype = next(self.pipe.vae.parameters()).dtype
        model_device = next(self.pipe.vae.parameters()).device
        
        print(f"模型设备: {model_device}")
        print(f"模型数据类型: {model_dtype}")
        
        # 将latents移动到正确的设备和数据类型
        latents = latents.to(device=model_device, dtype=model_dtype)
        
        print(f"解码latents形状: {latents.shape}")
        print(f"解码latents设备: {latents.device}")
        print(f"解码latents数据类型: {latents.dtype}")
        
        # 🔧 强制设置pipeline设备，确保所有操作在同一设备上
        self.pipe.device = model_device
        
        # 使用VAE解码
        with torch.no_grad():
            try:
                if tiled:
                    print("🔧 尝试tiled解码...")
                    decoded_video = self.pipe.decode_video(
                        latents, 
                        tiled=True, 
                        tile_size=tile_size, 
                        tile_stride=tile_stride
                    )
                else:
                    print("🔧 使用非tiled解码...")
                    decoded_video = self.pipe.decode_video(latents, tiled=False)
                    
            except Exception as e:
                print(f"decode_video失败，错误: {e}")
                import traceback
                traceback.print_exc()
                
                # 🔧 fallback: 尝试直接调用VAE
                try:
                    print("🔧 尝试直接调用VAE解码...")
                    decoded_video = self.pipe.vae.decode(
                        latents.squeeze(0),  # 移除batch维度 [C, T, H, W]
                        device=model_device, 
                        tiled=False
                    )
                    # 手动调整维度: VAE输出 [T, H, W, C] -> [1, T, H, W, C]
                    if len(decoded_video.shape) == 4:  # [T, H, W, C]
                        decoded_video = decoded_video.unsqueeze(0)  # -> [1, T, H, W, C]
                except Exception as e2:
                    print(f"直接VAE解码也失败: {e2}")
                    raise e2
        
        print(f"解码后视频形状: {decoded_video.shape}")
        
        # 🔧 关键修正：正确处理维度顺序
        video_np = None
        
        if len(decoded_video.shape) == 5:
            # 检查不同的可能维度顺序
            if decoded_video.shape == torch.Size([1, 3, 113, 480, 832]):
                # 格式: [B, C, T, H, W] -> 需要转换为 [T, H, W, C]
                print("🔧 检测到格式: [B, C, T, H, W]")
                video_np = decoded_video[0].permute(1, 2, 3, 0).to(torch.float32).cpu().numpy()  # [T, H, W, C]
            elif decoded_video.shape[1] == 3:
                # 如果第二个维度是3，可能是 [B, C, T, H, W]
                print("🔧 检测到可能的格式: [B, C, T, H, W]")
                video_np = decoded_video[0].permute(1, 2, 3, 0).to(torch.float32).cpu().numpy()  # [T, H, W, C]
            elif decoded_video.shape[-1] == 3:
                # 如果最后一个维度是3，可能是 [B, T, H, W, C]
                print("🔧 检测到格式: [B, T, H, W, C]")
                video_np = decoded_video[0].to(torch.float32).cpu().numpy()  # [T, H, W, C]
            else:
                # 尝试找到维度为3的位置
                shape = list(decoded_video.shape)
                if 3 in shape:
                    channel_dim = shape.index(3)
                    print(f"🔧 检测到通道维度在位置: {channel_dim}")
                    
                    if channel_dim == 1:  # [B, C, T, H, W]
                        video_np = decoded_video[0].permute(1, 2, 3, 0).to(torch.float32).cpu().numpy()
                    elif channel_dim == 4:  # [B, T, H, W, C]
                        video_np = decoded_video[0].to(torch.float32).cpu().numpy()
                    else:
                        print(f"⚠️ 未知的通道维度位置: {channel_dim}")
                        raise ValueError(f"Cannot handle channel dimension at position {channel_dim}")
                else:
                    print(f"⚠️ 未找到通道维度为3的位置，形状: {decoded_video.shape}")
                    raise ValueError(f"Cannot find channel dimension of size 3 in shape {decoded_video.shape}")
                    
        elif len(decoded_video.shape) == 4:
            # 4维张量，检查可能的格式
            if decoded_video.shape[-1] == 3:  # [T, H, W, C]
                video_np = decoded_video.to(torch.float32).cpu().numpy()
            elif decoded_video.shape[0] == 3:  # [C, T, H, W]
                video_np = decoded_video.permute(1, 2, 3, 0).to(torch.float32).cpu().numpy()
            else:
                print(f"⚠️ 无法处理的4D视频形状: {decoded_video.shape}")
                raise ValueError(f"Cannot handle 4D video tensor shape: {decoded_video.shape}")
        else:
            print(f"⚠️ 意外的视频维度数: {len(decoded_video.shape)}")
            raise ValueError(f"Unexpected video tensor dimensions: {decoded_video.shape}")
        
        if video_np is None:
            raise ValueError("Failed to convert video tensor to numpy array")
            
        print(f"转换后视频数组形状: {video_np.shape}")
        
        # 🔧 验证最终形状
        if len(video_np.shape) != 4:
            raise ValueError(f"Expected 4D array [T, H, W, C], got {video_np.shape}")
        
        if video_np.shape[-1] != 3:
            print(f"⚠️ 通道数异常: 期望3，实际{video_np.shape[-1]}")
            print(f"完整形状: {video_np.shape}")
            # 尝试其他维度排列
            if video_np.shape[0] == 3:  # [C, T, H, W]
                print("🔧 尝试重新排列: [C, T, H, W] -> [T, H, W, C]")
                video_np = np.transpose(video_np, (1, 2, 3, 0))
            elif video_np.shape[1] == 3:  # [T, C, H, W]
                print("🔧 尝试重新排列: [T, C, H, W] -> [T, H, W, C]")
                video_np = np.transpose(video_np, (0, 2, 3, 1))
            else:
                raise ValueError(f"Expected 3 channels (RGB), got {video_np.shape[-1]} channels")
        
        # 反归一化
        video_np = (video_np * 0.5 + 0.5).clip(0, 1)  # 反归一化
        video_np = (video_np * 255).astype(np.uint8)
        
        print(f"最终视频数组形状: {video_np.shape}")
        print(f"视频数组值范围: {video_np.min()} - {video_np.max()}")
        
        # 保存视频
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        try:
            with imageio.get_writer(output_path, fps=10, quality=8) as writer:
                for frame_idx, frame in enumerate(video_np):
                    # 🔧 验证每一帧的形状
                    if len(frame.shape) != 3 or frame.shape[-1] != 3:
                        print(f"⚠️ 帧 {frame_idx} 形状异常: {frame.shape}")
                        continue
                    
                    writer.append_data(frame)
                    if frame_idx % 10 == 0:
                        print(f"  写入帧 {frame_idx}/{len(video_np)}")
        except Exception as e:
            print(f"保存视频失败: {e}")
            # 🔧 尝试保存前几帧为图片进行调试
            debug_dir = os.path.join(os.path.dirname(output_path), "debug_frames")
            os.makedirs(debug_dir, exist_ok=True)
            
            for i in range(min(5, len(video_np))):
                frame = video_np[i]
                debug_path = os.path.join(debug_dir, f"debug_frame_{i}.png")
                try:
                    if len(frame.shape) == 3 and frame.shape[-1] == 3:
                        Image.fromarray(frame).save(debug_path)
                        print(f"调试: 保存帧 {i} 到 {debug_path}")
                    else:
                        print(f"调试: 帧 {i} 形状异常: {frame.shape}")
                except Exception as e2:
                    print(f"调试: 保存帧 {i} 失败: {e2}")
            raise e
        
        print(f"✅ 视频保存到: {output_path}")
        return video_np

    def save_frames_as_images(self, video_np, output_dir, prefix="frame"):
        """将视频帧保存为单独的图像文件"""
        os.makedirs(output_dir, exist_ok=True)
        
        for i, frame in enumerate(video_np):
            frame_path = os.path.join(output_dir, f"{prefix}_{i:04d}.png")
            # 🔧 验证帧形状
            if len(frame.shape) == 3 and frame.shape[-1] == 3:
                Image.fromarray(frame).save(frame_path)
            else:
                print(f"⚠️ 跳过形状异常的帧 {i}: {frame.shape}")
        
        print(f"✅ 保存了 {len(video_np)} 帧到: {output_dir}")

def decode_single_episode(encoded_pth_path, vae_path, output_base_dir, device="cuda"):
    """解码单个episode的编码数据 - 修正版本"""
    print(f"\n🔧 解码episode: {encoded_pth_path}")
    
    # 加载编码数据
    try:
        encoded_data = torch.load(encoded_pth_path, weights_only=False, map_location="cpu")
        print(f"✅ 成功加载编码数据")
    except Exception as e:
        print(f"❌ 加载编码数据失败: {e}")
        return False
    
    # 检查数据结构
    print("🔍 编码数据结构:")
    for key, value in encoded_data.items():
        if isinstance(value, torch.Tensor):
            print(f"  - {key}: {value.shape}, dtype: {value.dtype}, device: {value.device}")
        elif isinstance(value, dict):
            print(f"  - {key}: dict with keys {list(value.keys())}")
        else:
            print(f"  - {key}: {type(value)}")
    
    # 获取latents
    latents = encoded_data.get('latents')
    if latents is None:
        print("❌ 未找到latents数据")
        return False
    
    # 🔧 确保latents在CPU上（加载时的默认状态）
    if latents.device != torch.device('cpu'):
        latents = latents.cpu()
        print(f"🔧 将latents移动到CPU: {latents.device}")
    
    episode_info = encoded_data.get('episode_info', {})
    episode_idx = episode_info.get('episode_idx', 'unknown')
    total_frames = episode_info.get('total_frames', latents.shape[1] * 4)  # 估算原始帧数
    
    print(f"Episode信息:")
    print(f"  - Episode索引: {episode_idx}")
    print(f"  - Latents形状: {latents.shape}")
    print(f"  - Latents设备: {latents.device}")
    print(f"  - Latents数据类型: {latents.dtype}")
    print(f"  - 原始总帧数: {total_frames}")
    print(f"  - 压缩后帧数: {latents.shape[1]}")
    
    # 创建输出目录
    episode_name = f"episode_{episode_idx:06d}" if isinstance(episode_idx, int) else f"episode_{episode_idx}"
    output_dir = os.path.join(output_base_dir, episode_name)
    os.makedirs(output_dir, exist_ok=True)
    
    # 初始化解码器
    try:
        decoder = VideoDecoder(vae_path, device)
    except Exception as e:
        print(f"❌ 初始化解码器失败: {e}")
        return False
    
    # 解码为视频
    video_output_path = os.path.join(output_dir, "decoded_video.mp4")
    try:
        video_np = decoder.decode_latents_to_video(
            latents, 
            video_output_path,
            tiled=False,  # 🔧 首先尝试非tiled解码，避免tiled的复杂性
            tile_size=(34, 34),
            tile_stride=(18, 16)
        )
        
        # 保存前几帧为图像（用于快速检查）
        frames_dir = os.path.join(output_dir, "frames")
        sample_frames = video_np[:min(10, len(video_np))]  # 只保存前10帧
        decoder.save_frames_as_images(sample_frames, frames_dir, f"frame_{episode_idx}")
        
        # 保存解码信息
        decode_info = {
            "source_pth": encoded_pth_path,
            "decoded_video_path": video_output_path,
            "latents_shape": list(latents.shape),
            "decoded_video_shape": list(video_np.shape),
            "original_total_frames": total_frames,
            "decoded_frames": len(video_np),
            "compression_ratio": total_frames / len(video_np) if len(video_np) > 0 else 0,
            "latents_dtype": str(latents.dtype),
            "latents_device": str(latents.device),
            "vae_compression_ratio": total_frames / latents.shape[1] if latents.shape[1] > 0 else 0
        }
        
        info_path = os.path.join(output_dir, "decode_info.json")
        with open(info_path, 'w') as f:
            json.dump(decode_info, f, indent=2)
        
        print(f"✅ Episode {episode_idx} 解码完成")
        print(f"  - 原始帧数: {total_frames}")
        print(f"  - 解码帧数: {len(video_np)}")
        print(f"  - 压缩比: {decode_info['compression_ratio']:.2f}")
        print(f"  - VAE时间压缩比: {decode_info['vae_compression_ratio']:.2f}")
        return True
        
    except Exception as e:
        print(f"❌ 解码失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def batch_decode_episodes(encoded_base_dir, vae_path, output_base_dir, max_episodes=None, device="cuda"):
    """批量解码episodes"""
    print(f"🔧 批量解码Open-X episodes")
    print(f"源目录: {encoded_base_dir}")
    print(f"输出目录: {output_base_dir}")
    
    # 查找所有编码的episodes
    episode_dirs = []
    if os.path.exists(encoded_base_dir):
        for item in sorted(os.listdir(encoded_base_dir)):  # 排序确保一致性
            episode_dir = os.path.join(encoded_base_dir, item)
            if os.path.isdir(episode_dir):
                encoded_path = os.path.join(episode_dir, "encoded_video.pth")
                if os.path.exists(encoded_path):
                    episode_dirs.append(encoded_path)
    
    print(f"找到 {len(episode_dirs)} 个编码的episodes")
    
    if max_episodes and len(episode_dirs) > max_episodes:
        episode_dirs = episode_dirs[:max_episodes]
        print(f"限制处理前 {max_episodes} 个episodes")
    
    # 批量解码
    success_count = 0
    for i, encoded_pth_path in enumerate(tqdm(episode_dirs, desc="解码episodes")):
        print(f"\n{'='*60}")
        print(f"处理 {i+1}/{len(episode_dirs)}: {os.path.basename(os.path.dirname(encoded_pth_path))}")
        
        success = decode_single_episode(encoded_pth_path, vae_path, output_base_dir, device)
        if success:
            success_count += 1
        
        print(f"当前成功率: {success_count}/{i+1} ({success_count/(i+1)*100:.1f}%)")
    
    print(f"\n🎉 批量解码完成!")
    print(f"总处理: {len(episode_dirs)} 个episodes")
    print(f"成功解码: {success_count} 个episodes")
    print(f"成功率: {success_count/len(episode_dirs)*100:.1f}%")

def main():
    parser = argparse.ArgumentParser(description="解码Open-X编码的latents以验证正确性 - 修正版本")
    parser.add_argument("--mode", type=str, choices=["single", "batch"], default="batch",
                       help="解码模式：single (单个episode) 或 batch (批量)")
    parser.add_argument("--encoded_pth", type=str, 
                       default="/share_zhuyixuan05/zhuyixuan05/openx-fractal-encoded/episode_000000/encoded_video.pth",
                       help="单个编码文件路径（single模式）")
    parser.add_argument("--encoded_base_dir", type=str,
                       default="/share_zhuyixuan05/zhuyixuan05/openx-fractal-encoded",
                       help="编码数据基础目录（batch模式）")
    parser.add_argument("--vae_path", type=str,
                       default="models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
                       help="VAE模型路径")
    parser.add_argument("--output_dir", type=str,
                       default="./decoded_results_fixed",
                       help="解码输出目录")
    parser.add_argument("--max_episodes", type=int, default=5,
                       help="最大解码episodes数量（batch模式，用于测试）")
    parser.add_argument("--device", type=str, default="cuda",
                       help="计算设备")
    
    args = parser.parse_args()
    
    print("🔧 Open-X Latents 解码验证工具 (修正版本 - Fixed)")
    print(f"模式: {args.mode}")
    print(f"VAE路径: {args.vae_path}")
    print(f"输出目录: {args.output_dir}")
    print(f"设备: {args.device}")
    
    # 🔧 检查CUDA可用性
    if args.device == "cuda" and not torch.cuda.is_available():
        print("⚠️ CUDA不可用，切换到CPU")
        args.device = "cpu"
    
    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)
    
    if args.mode == "single":
        print(f"输入文件: {args.encoded_pth}")
        if not os.path.exists(args.encoded_pth):
            print(f"❌ 输入文件不存在: {args.encoded_pth}")
            return
        
        success = decode_single_episode(args.encoded_pth, args.vae_path, args.output_dir, args.device)
        if success:
            print("✅ 单个episode解码成功")
        else:
            print("❌ 单个episode解码失败")
    
    elif args.mode == "batch":
        print(f"输入目录: {args.encoded_base_dir}")
        print(f"最大episodes: {args.max_episodes}")
        
        if not os.path.exists(args.encoded_base_dir):
            print(f"❌ 输入目录不存在: {args.encoded_base_dir}")
            return
        
        batch_decode_episodes(args.encoded_base_dir, args.vae_path, args.output_dir, args.max_episodes, args.device)

if __name__ == "__main__":
    main()