pruned_rebuild.py

#!/usr/bin/env python3
"""
从 safetensors + config.json 精确重建剪枝后的 UNet 结构。

关键思路：
  不使用 align_tensor 填充零值（会污染已学习的权重）。
  而是先把标准 UNet 里的每个 Conv2d/Linear 替换为 safetensors
  中实际形状对应的新模块，再用 load_state_dict 加载。
"""

import os
import sys
import json
import torch
import torch.nn as nn
from safetensors.torch import load_file

sys.path.insert(0, '/home/ubuntu')

os.environ.update({
    'HF_HOME': '/opt/dlami/nvme/hf_cache',
    'TRANSFORMERS_CACHE': '/opt/dlami/nvme/hf_cache',
    'TMPDIR': '/opt/dlami/nvme/tmp'
})


# ---------------------------------------------------------------------------
# 核心：从 safetensors + config.json 重建剪枝模型
# ---------------------------------------------------------------------------

def _get_parent_and_attr(model: nn.Module, dotted_name: str):
    """返回 (parent_module, attr_name)，用于 setattr 替换子模块。"""
    parts = dotted_name.split('.')
    obj = model
    for p in parts[:-1]:
        obj = getattr(obj, p)
    return obj, parts[-1]


def _find_num_groups(original_num_groups: int, new_num_channels: int) -> int:
    """找到能整除 new_num_channels 的最大 num_groups（不超过 original_num_groups）。"""
    ng = original_num_groups
    while ng > 1:
        if new_num_channels % ng == 0:
            return ng
        ng //= 2
    return 1


def _replace_layers_to_match_shapes(unet: nn.Module, st: dict) -> int:
    """
    遍历 unet 所有 Conv2d / Linear / GroupNorm，
    若 safetensors 中对应权重形状不同，则替换为正确尺寸的新模块。
    返回替换的层数量。
    """
    replaced = 0
    for name, module in list(unet.named_modules()):
        weight_key = name + '.weight'
        if weight_key not in st:
            continue

        w = st[weight_key]
        has_bias = (name + '.bias') in st

        if isinstance(module, nn.Conv2d):
            out_c, in_c = w.shape[0], w.shape[1]
            if out_c != module.out_channels or in_c != module.in_channels:
                new_mod = nn.Conv2d(
                    in_c, out_c,
                    kernel_size=module.kernel_size,
                    stride=module.stride,
                    padding=module.padding,
                    dilation=module.dilation,
                    groups=module.groups,
                    bias=has_bias,
                )
                parent, attr = _get_parent_and_attr(unet, name)
                setattr(parent, attr, new_mod)
                replaced += 1

        elif isinstance(module, nn.Linear):
            out_f, in_f = w.shape[0], w.shape[1]
            if out_f != module.out_features or in_f != module.in_features:
                new_mod = nn.Linear(in_f, out_f, bias=has_bias)
                parent, attr = _get_parent_and_attr(unet, name)
                setattr(parent, attr, new_mod)
                replaced += 1

        elif isinstance(module, nn.GroupNorm):
            new_num_ch = w.shape[0]
            if new_num_ch != module.num_channels:
                ng = _find_num_groups(module.num_groups, new_num_ch)
                new_mod = nn.GroupNorm(ng, new_num_ch, eps=module.eps, affine=module.affine)
                parent, attr = _get_parent_and_attr(unet, name)
                setattr(parent, attr, new_mod)
                replaced += 1

        elif isinstance(module, nn.LayerNorm):
            # transformer_blocks.*.norm1/2/3 使用 LayerNorm，normalized_shape=[dim]
            new_dim = w.shape[0]
            if list(module.normalized_shape) != [new_dim]:
                new_mod = nn.LayerNorm(new_dim, eps=module.eps, elementwise_affine=module.elementwise_affine)
                parent, attr = _get_parent_and_attr(unet, name)
                setattr(parent, attr, new_mod)
                replaced += 1

    return replaced


def _fix_internal_attrs(unet: nn.Module):
    """
    更新 diffusers UNet 内部依赖于通道数的属性
    （Upsample2D.channels、ResnetBlock2D.in_channels 等）。
    """
    for name, module in unet.named_modules():
        if hasattr(module, 'channels') and hasattr(module, 'conv'):
            if hasattr(module.conv, 'in_channels'):
                module.channels = module.conv.in_channels
        if hasattr(module, 'in_channels') and hasattr(module, 'conv1'):
            if hasattr(module.conv1, 'in_channels'):
                module.in_channels = module.conv1.in_channels
        if hasattr(module, 'out_channels') and hasattr(module, 'conv2'):
            if hasattr(module.conv2, 'out_channels'):
                module.out_channels = module.conv2.out_channels
        if hasattr(module, 'to_q') and hasattr(module, 'inner_dim'):
            if hasattr(module.to_q, 'weight'):
                new_inner_dim = module.to_q.weight.shape[0]
                old_inner_dim = module.inner_dim
                module.inner_dim = new_inner_dim
                if hasattr(module, 'inner_kv_dim'):
                    module.inner_kv_dim = new_inner_dim
                # Update heads: head_dim is invariant, recompute heads count
                if hasattr(module, 'heads') and module.heads > 0 and old_inner_dim > 0:
                    head_dim = old_inner_dim // module.heads
                    if head_dim > 0 and new_inner_dim % head_dim == 0:
                        module.heads = new_inner_dim // head_dim
                        if hasattr(module, 'sliceable_head_dim'):
                            module.sliceable_head_dim = module.heads


def create_unet_from_safetensors(safetensors_path: str, config_path: str = None) -> nn.Module:
    """
    从 safetensors + config.json 精确重建剪枝后的 UNet。

    流程：
      1. 加载 safetensors（获取实际张量形状）
      2. 从 config_path 中的 model_config 构建标准 UNet
      3. 将形状不匹配的 Conv2d/Linear 替换为正确尺寸
      4. load_state_dict
      5. 修复内部属性
    """
    from diffusers import UNet2DConditionModel

    # 1. 加载 safetensors
    print(f"加载 safetensors: {safetensors_path}")
    st = load_file(safetensors_path)
    total_params = sum(v.numel() for v in st.values())
    print(f"  safetensors 共 {len(st)} 个张量，{total_params/1e6:.1f}M 参数")

    # 2. 读取 model_config
    if config_path is None:
        config_path = safetensors_path.replace('.safetensors', '.config.json')

    model_config = None
    if os.path.exists(config_path):
        with open(config_path, 'r', encoding='utf-8') as f:
            meta = json.load(f)
        model_config = meta.get('model_config')
        print(f"  读取配置: {config_path}")

    # 回退到默认 SD 1.5 配置
    if not model_config or not isinstance(model_config, dict):
        print("  ⚠️  未找到 model_config，使用 SD 1.5 默认配置")
        model_config = {
            "sample_size": 64,
            "in_channels": 4,
            "out_channels": 4,
            "layers_per_block": 2,
            "block_out_channels": [320, 640, 1280, 1280],
            "down_block_types": [
                "CrossAttnDownBlock2D", "CrossAttnDownBlock2D",
                "CrossAttnDownBlock2D", "DownBlock2D"
            ],
            "up_block_types": [
                "UpBlock2D", "CrossAttnUpBlock2D",
                "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"
            ],
            "cross_attention_dim": 768,
            "attention_head_dim": 8,
        }

    # 3. 构建标准 UNet
    print("  构建标准 UNet 架构...")
    unet = UNet2DConditionModel(**model_config)

    # 4. 将形状不匹配的层替换为正确尺寸
    replaced = _replace_layers_to_match_shapes(unet, st)
    print(f"  替换了 {replaced} 个形状不匹配的层")

    # 5. 加载 safetensors 权重
    missing, unexpected = unet.load_state_dict(st, strict=False)
    if missing:
        print(f"  ⚠️  缺失键: {len(missing)} 个（例如 {missing[:3]}）")
    if unexpected:
        print(f"  ⚠️  多余键: {len(unexpected)} 个")

    # 6. 修复内部属性
    _fix_internal_attrs(unet)

    param_count = sum(p.numel() for p in unet.parameters())
    print(f"  ✅ 重建完成，参数量: {param_count/1e6:.1f}M")
    return unet


# ---------------------------------------------------------------------------
# 验证：前向推理测试
# ---------------------------------------------------------------------------

def verify_forward(unet: nn.Module, device: str = 'cpu') -> bool:
    """对重建的模型跑一次前向推理，验证输出形状正确。"""
    unet = unet.to(device).eval()
    with torch.no_grad():
        sample = torch.randn(1, 4, 64, 64, device=device)
        timestep = torch.tensor([1], device=device)
        enc_hs = torch.randn(1, 77, 768, device=device)
        try:
            out = unet(sample, timestep, encoder_hidden_states=enc_hs)
            assert tuple(out.sample.shape) == (1, 4, 64, 64), \
                f"输出形状异常: {out.sample.shape}"
            print(f"  前向推理 OK，输出形状: {tuple(out.sample.shape)}")
            return True
        except Exception as e:
            print(f"  ❌ 前向推理失败: {e}")
            import traceback
            traceback.print_exc()
            return False


# ---------------------------------------------------------------------------
# 主入口
# ---------------------------------------------------------------------------

def main():
    safetensors_path = os.environ.get(
        'PRUNED_SAFETENS_PATH',
        '/opt/dlami/nvme/prune_outputs/taylor_sp_unet_v2.safetensors'
    )
    config_path = safetensors_path.replace('.safetensors', '.config.json')

    print("=" * 60)
    print("从 safetensors + config.json 重建剪枝 UNet")
    print("=" * 60)

    unet = create_unet_from_safetensors(safetensors_path, config_path)
    ok = verify_forward(unet)

    if ok:
        print("\n✅ 模型重建成功，可直接用于推理/蒸馏！")
    else:
        print("\n❌ 模型重建后前向推理失败，请检查配置")


if __name__ == '__main__':
    main()