init
This commit is contained in:
smthemex
+195
@@ -0,0 +1,195 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
import torch
|
||||
import gc
|
||||
import comfy.model_management as mm
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from comfy.utils import common_upscale
|
||||
from datetime import datetime
|
||||
import folder_paths
|
||||
import soundfile as sf
|
||||
import uuid
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from transformers import AutoConfig
|
||||
from safetensors.torch import load_file
|
||||
from .SheetSage2.modeling_sheetsage2 import SheetSage2Model
|
||||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def build_mert2_encoder(config, files):
|
||||
from .SheetSage2.configuration_mert2 import MERT2Config
|
||||
from .SheetSage2.modeling_mert2 import MERT2Model
|
||||
from .SheetSage2.modeling_sheetsage2 import BASE_CODE_HASHES
|
||||
encoder_config = MERT2Config(**config.backbone_config)
|
||||
encoder_config._attn_implementation = config.encoder_attn_implementation
|
||||
encoder = MERT2Model(encoder_config)
|
||||
expected_hash = getattr(config, "base_model_sha256", None)
|
||||
state = load_file(files)
|
||||
keys = set(encoder.state_dict())
|
||||
if not keys & set(state):
|
||||
for prefix in ("model.", "encoder."):
|
||||
stripped = {key[len(prefix):] if key.startswith(prefix) else key: value for key, value in state.items()}
|
||||
if keys & set(stripped):
|
||||
state = stripped
|
||||
break
|
||||
report = encoder.load_state_dict(state, strict=False)
|
||||
del state
|
||||
if report.missing_keys:
|
||||
raise ValueError(f"MERT-v2 parent 权重不全，缺 {len(report.missing_keys)} 个键: {report.missing_keys[:8]}")
|
||||
return encoder.to(dtype=torch.float32)
|
||||
|
||||
|
||||
def loadsheetsage(model_file_path, mert2_file, dtype=torch.float32, device=torch.device("cpu")):
|
||||
"""序列化加载 SheetSage2，不调用 from_pretrained。
|
||||
|
||||
- 权重含 encoder.*（merged 快照）: 直接按 merged 配置构建后灌权重，无需父模型。
|
||||
- 权重只有 adapter.*（原始发布版）: 建 encoder 骨架 + 灌父权重 + merge_lora（必须在 CPU/float32 上合并）。
|
||||
"""
|
||||
config = AutoConfig.from_pretrained(os.path.join(cur_path, "SheetSage2"), trust_remote_code=True, local_files_only=True)
|
||||
state = load_file(model_file_path)
|
||||
if any(key.startswith("encoder.") for key in state):
|
||||
config.weights_format = "merged" # 必须在构造之前改：__init__ 靠它决定建不建 encoder
|
||||
model = SheetSage2Model(config)
|
||||
model.load_state_dict(state, strict=False)
|
||||
else:
|
||||
model = SheetSage2Model(config) # adapter 分支: encoder=None, adapter=EncoderAdapters, lora_merged=False
|
||||
model.load_state_dict(state, strict=False)
|
||||
missing = [key for key in model.state_dict() if key.startswith("adapter.") and key not in state]
|
||||
if missing:
|
||||
raise ValueError(f"Adapter 权重不全，缺 {len(missing)} 个键: {missing[:8]}")
|
||||
model.encoder = build_mert2_encoder(config, mert2_file)
|
||||
model.merge_lora() # 把 adapter 的 LoRA 加进主干，并置 lora_merged=True
|
||||
del state
|
||||
model.to(dtype=torch.float32) # 合并/加载统一在 float32
|
||||
if dtype is not None:
|
||||
model.to(dtype=dtype)
|
||||
if device is not None:
|
||||
model.to(device)
|
||||
return model.eval().requires_grad_(False)
|
||||
|
||||
|
||||
def load_yue2_pipeline(repo="m-a-p/YuE2-3B",yue_ckpt=None,vae=None,device="cuda"):
|
||||
from .src.yue2.pipeline import YuE2Pipeline
|
||||
return YuE2Pipeline.from_pretrained(repo, yue_ckpt=yue_ckpt, vae=vae, device=device)
|
||||
|
||||
def yue2_generate(pipe,request,vae_model=None):
|
||||
song = pipe(**request,vae_model=vae_model)
|
||||
return song
|
||||
|
||||
|
||||
def audio2path(audio,):
|
||||
prefix=datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||||
audio_file = os.path.join(folder_paths.get_temp_directory(), f"audio_refer_temp_{prefix}.wav")
|
||||
waveform = audio["waveform"].squeeze(0)
|
||||
waveform_np = waveform.cpu().numpy() if hasattr(waveform, 'cpu') else waveform.numpy()
|
||||
|
||||
# 3. 格式转换：torchaudio 格式为 (channels, samples)，soundfile 需要 (samples, channels)
|
||||
# 如果是单声道音频 (1, samples)，转置后变成 (samples, 1)，符合 soundfile 的单声道要求
|
||||
if waveform_np.ndim == 2:
|
||||
waveform_np = waveform_np.T
|
||||
|
||||
sf.write(audio_file, waveform_np, audio["sample_rate"])
|
||||
|
||||
return audio_file
|
||||
|
||||
|
||||
def clear_comfyui_cache():
|
||||
cf_models=mm.loaded_models()
|
||||
try:
|
||||
for pipe in cf_models:
|
||||
pipe.unpatch_model(device_to=torch.device("cpu"))
|
||||
except: pass
|
||||
mm.soft_empty_cache()
|
||||
torch.cuda.empty_cache()
|
||||
max_gpu_memory = torch.cuda.max_memory_allocated()
|
||||
print(f"After Max GPU memory allocated: {max_gpu_memory / 1000 ** 3:.2f} GB")
|
||||
|
||||
def gc_cleanup():
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def phi2narry(img):
|
||||
img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return img
|
||||
|
||||
def tensor2image(tensor):
|
||||
tensor = tensor.cpu()
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def tensor2pillist(tensor_in):
|
||||
d1, _, _, _ = tensor_in.size()
|
||||
if d1 == 1:
|
||||
img_list = [tensor2image(tensor_in)]
|
||||
else:
|
||||
tensor_list = torch.chunk(tensor_in, chunks=d1)
|
||||
img_list=[tensor2image(i) for i in tensor_list]
|
||||
return img_list
|
||||
|
||||
def tensor2pillist_upscale(tensor_in,width,height):
|
||||
d1, _, _, _ = tensor_in.size()
|
||||
if d1 == 1:
|
||||
img_list = [nomarl_upscale(tensor_in,width,height)]
|
||||
else:
|
||||
tensor_list = torch.chunk(tensor_in, chunks=d1)
|
||||
img_list=[nomarl_upscale(i,width,height) for i in tensor_list]
|
||||
return img_list
|
||||
|
||||
def tensor2list(tensor_in,width,height):
|
||||
if tensor_in is None:
|
||||
return None
|
||||
d1, _, _, _ = tensor_in.size()
|
||||
if d1 == 1:
|
||||
tensor_list = [tensor_upscale(tensor_in,width,height)]
|
||||
else:
|
||||
tensor_list_ = torch.chunk(tensor_in, chunks=d1)
|
||||
tensor_list=[tensor_upscale(i,width,height) for i in tensor_list_]
|
||||
return tensor_list
|
||||
|
||||
def tensor_upscale(tensor, width, height):
|
||||
samples = tensor.movedim(-1, 1)
|
||||
samples = common_upscale(samples, width, height, "bilinear", "center")
|
||||
samples = samples.movedim(1, -1)
|
||||
return samples
|
||||
|
||||
def nomarl_upscale(img, width, height):
|
||||
samples = img.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "bilinear", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
img = tensor2image(samples)
|
||||
return img
|
||||
|
||||
|
||||
def map_0_1_to_neg1_1(t):
|
||||
|
||||
if not torch.is_tensor(t):
|
||||
t = torch.tensor(t)
|
||||
t = t.float()
|
||||
|
||||
try:
|
||||
vmax = float(t.max())
|
||||
except Exception:
|
||||
vmax = 1.0
|
||||
if vmax > 2.0:
|
||||
t = t / 255.0
|
||||
try:
|
||||
vmin = float(t.min())
|
||||
vmax = float(t.max())
|
||||
except Exception:
|
||||
vmin, vmax = -1.0, 1.0
|
||||
if vmin >= 0.0 and vmax <= 1.1:
|
||||
t = t * 2.0 - 1.0
|
||||
return t
|
||||
|
||||
def map_neg1_1_to_0_1(t):
|
||||
if not torch.is_tensor(t):
|
||||
t = torch.tensor(t)
|
||||
t = t.float()
|
||||
t = (t + 1.0) * 0.5
|
||||
t = t.clamp(0.0, 1.0)
|
||||
return t
|
||||
Reference in New Issue
Block a user