scripts/cubemap_vae.py

import torch
from diffusers.models import AutoencoderKL
from torch import nn
from torchvision.transforms import ToPILImage
from torch import Tensor

class SynchronizedGroupNorm(nn.Module):
    def __init__(self, original_group_norm: nn.GroupNorm, num_views: int = 6):
        super().__init__()
        self.num_views = num_views
        
        # 继承原始分组参数
        self.num_groups = original_group_norm.num_groups
        self.num_channels = original_group_norm.num_channels
        self.eps = original_group_norm.eps
        
        # 按照通道组重构参数
        self.group_size = self.num_channels // self.num_groups
        
        # 继承原始参数（保持每个分组的仿射变换）
        self.weight = nn.Parameter(original_group_norm.weight.detach().view(self.num_groups, self.group_size))
        self.bias = nn.Parameter(original_group_norm.bias.detach().view(self.num_groups, self.group_size))
    
    def forward(self, x: torch.Tensor):
        """ 兼容 3D (B, C, D) 和 4D (B, C, H, W) 输入 """
        
        #print(f"Input shape: {x.shape}")  # Debugging

        # 获取输入的形状信息
        BxT, C = x.shape[:2]  # 只取前两个维度
        B = BxT // self.num_views  # 计算 batch 维度

        # 处理 3D (B, C, D) 输入
        if x.dim() == 3:
            D = x.shape[2]
            x = x.view(B, self.num_views, self.num_groups, self.group_size, D)

            # 计算 GroupNorm
            mean = x.mean(dim=(1, 3, 4), keepdim=True)
            var = x.var(dim=(1, 3, 4), keepdim=True, unbiased=False)
            x = (x - mean) / torch.sqrt(var + self.eps)
            # **修正 weight 和 bias 的形状**
            weight = self.weight.view(1, self.num_groups, self.group_size, 1)
            bias = self.bias.view(1, self.num_groups, self.group_size, 1)
            x = x * weight + bias

            # 还原形状
            return x.view(BxT, C, D)

        # 处理 4D (B, C, H, W) 输入
        elif x.dim() == 4:
            H, W = x.shape[2:]
            x = x.view(B, self.num_views, self.num_groups, self.group_size, H, W)

            # 计算 GroupNorm
            mean = x.mean(dim=(1, 3, 4, 5), keepdim=True)
            var = x.var(dim=(1, 3, 4, 5), keepdim=True, unbiased=False)
            x = (x - mean) / torch.sqrt(var + self.eps)
            # **修正 weight 和 bias 的形状**
            weight = self.weight.view(1, self.num_groups, self.group_size, 1, 1)
            bias = self.bias.view(1, self.num_groups, self.group_size, 1, 1)
            x = x * weight + bias

            # 还原形状
            return x.view(BxT, C, H, W)

        else:
            raise ValueError(f"Unsupported input shape: {x.shape}, expected 3D (B, C, D) or 4D (B, C, H, W).")
        

class CubemapVAE(AutoencoderKL):
    def __init__(self, pretrained_vae, num_views=6, in_channels=3,image_size=512):
        super().__init__(  # 继承自 AutoencoderKL
            act_fn="silu",
            block_out_channels=[128, 256, 512, 512],
            down_block_types=[
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "DownEncoderBlock2D"
            ],
            up_block_types=[
                "UpDecoderBlock2D",
                "UpDecoderBlock2D",
                "UpDecoderBlock2D",
                "UpDecoderBlock2D"
            ],
            latent_channels=pretrained_vae.config.latent_channels,
            in_channels=in_channels,  
            out_channels=in_channels  
        )
        self.num_views = num_views
        self.in_channels = in_channels
        
        
        # --- 替换关键模块，适配 Cubemap ---
        # 原 AutoencoderKL 的编码器不够灵活，直接覆盖编码器
        #self.encoder = CubemapEncoder(pretrained_encoder=pretrained_vae.encoder,num_views=num_views, in_channels=in_channels)
        #self.decoder = CubemapDecoder(pretrained_decoder=pretrained_vae.decoder, num_views=num_views, out_channels=in_channels,in_channels=4)
        self.encoder=pretrained_vae.encoder
        self.decoder=pretrained_vae.decoder
        self.quant_conv=pretrained_vae.quant_conv
        self.post_quant_conv=pretrained_vae.post_quant_conv
        # 将原 GroupNorm 替换为同步 GroupNorm
        replace_group_norm_with_sgn(self, num_views=num_views)
        
    def encode(self, images,return_dict:bool=True):
        batch_size, num_views, num_channels, height, width = images.shape
        images = images.view(batch_size*num_views,num_channels, height, width)
        return super().encode(images,return_dict=return_dict)
        
        
    def decode(self, latents, return_dict=True, **kwargs):
        """
        自定义 VAE 解码:
        - 去掉 UV 通道 (只保留前 4 个 latent 通道)
        - 调用原始 VAE 解码流程
        """
        
        print("Decoder Recieve Latent Shape:", latents.shape)
        # 确保 latents 至少有 4 个通道
        if latents.shape[1] > 4:
            latents = latents[:, :4, :, :]  # 只保留前 4 个通道，去掉 UV 通道
            
        
        
        return super().decode(latents, return_dict=return_dict, **kwargs)
    
    def decode_to_tensor(self, latents):
        decoded = self.decode(latents).sample  # (B*6, 3, H, W)
        
        B = latents.shape[0] // 6
        images = torch.split(decoded, B, dim=0)  # 按 batch 拆分
        
        return images  # Tuple of 6 tensors
    
    def decode_to_pil_images(self, latents:Tensor):
        images = self.decode_to_tensor(latents)  # 获取 6 张图
        to_pil = ToPILImage()
        
        return [to_pil(img[0].cpu().detach()) for img in images]  # 转换为 PIL



def replace_group_norm_with_sgn(model, num_views):
    """ 遍历 model，找到所有 GroupNorm 并替换成 SynchronizedGroupNorm """
    replacements = []  # 先收集要替换的 module 名称
    for name, module in model.named_modules():
        if isinstance(module, nn.GroupNorm):
            replacements.append(name)

    for name in replacements:
        parent_module, attr_name = get_parent_module(model, name)
        setattr(parent_module, attr_name, SynchronizedGroupNorm(getattr(parent_module, attr_name), num_views))

def get_parent_module(model, module_name):
    """ 获取 `module_name` 所在的上一级 module 和属性名称 """
    names = module_name.split(".")
    parent_module = model
    for name in names[:-1]:  # 遍历到倒数第二层
        parent_module = getattr(parent_module, name)
    return parent_module, names[-1]

            

def flatten_face_names(face_names):
    flat_face_names = []
    for item in face_names:
        if isinstance(item, str):  # 直接是字符串
            flat_face_names.append(item)
        elif isinstance(item, (list,tuple)):  # 是列表，展开其中的字符串
            flat_face_names.extend(item)
        else:
            raise ValueError(f"Unexpected type in face_names: {type(item)}")
    return flat_face_names