23 lines
575 B
Python
23 lines
575 B
Python
import torch
|
|
|
|
from ..modules.scaleu import ScaleU
|
|
|
|
|
|
def get_scaleu_patch(scaleu_nets):
|
|
def scaleu_patch(h, hsp, transformer_options):
|
|
_, idx = transformer_options['block']
|
|
sk = scaleu_nets[idx](h, hsp)
|
|
return sk
|
|
|
|
return scaleu_patch
|
|
|
|
|
|
def prepare_scaleu_nets(scaleu_ckpt) -> torch.nn.Module:
|
|
scaleu_nets = []
|
|
for i in range(12):
|
|
ckpt = scaleu_ckpt[f'{i}']
|
|
scaleu = ScaleU(True, len(ckpt['scaleu_b']), len(ckpt['scaleu_s']))
|
|
scaleu.load_state_dict(ckpt)
|
|
scaleu_nets.append(scaleu)
|
|
return scaleu_nets
|